# lora微调模型
import os
import torch
import torch.nn as nn
from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer
from datasets import load_dataset

ws = "code/try"
# os.environ["CUDA_VISIBLE_DEVICES"]="7"
if __name__ == "__main__":
    model_name = "glm-4-9b-chat-hf" 
    print(f"加载模型: {model_name}")
    model_path = os.path.join(ws, f"models/{model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_path,trust_remote_code=True)

    # 加载数据
    tra_dataset_path = "code/try/dataset/finetune/train_zh_0.json"
    tra_ds = load_dataset("json", data_files={"train": tra_dataset_path})['train'].select(range(50000))
    val_dataset_path = "code/try/dataset/finetune/valid_zh_0.json"
    val_ds = load_dataset("json", data_files={"val": val_dataset_path})['val'] #.select(range(5000))

    # 构造数据集
    def proc_func(examples):
        # 参考medicalGPT的实现，labels与input对齐，仅assistant部分计算loss；结尾处添加eos；tokenizer encode构造对话数据
        inputs = tokenizer.apply_chat_template([
            {"role":"user","content":examples["instruction"]+"\n"+examples["input"]},
            {"role":"assistant","content":examples["output"]},
        ],truncation=True,return_dict=False) # 
        # labels中将<|assistant|>之前的置-100，不予计算loss
        assistant_token = tokenizer.encode('<|assistant|>',add_special_tokens=False)[0]
        labels = [inputs[i] if i>=inputs.index(assistant_token) else -100 for i in range(len(inputs))]
        mask = [1]* len(inputs)
        return {"input_ids":inputs,"attention_mask":mask,"labels":labels}
    
    tra_ds = tra_ds.map(proc_func) #,batched=True
    val_ds = val_ds.map(proc_func) #,batched=True

    # glm4微调模型类型一定要为bfloat16，否则loss会出现nan的情况。
    model = AutoModelForCausalLM.from_pretrained(model_path,
                                                 trust_remote_code=True,
                                                 torch_dtype=torch.bfloat16,
                                                 ) #.eval()
    model.gradient_checkpointing_enable()
    
    # 训练配置
    args = Seq2SeqTrainingArguments( # 微调一定要用Seq2SeqTrainingArguments和Seq2SeqTrainer，否则loss就会为nan
        output_dir=os.path.join(ws, f"models/sft/{model_name}"),
        eval_strategy="epoch",
        # eval_steps=100,
        learning_rate=1e-5,
        num_train_epochs=3,
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        save_strategy="no",        
    )
    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        lora_dropout=0.1,
        peft_type="LORA", 
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj","v_proj"],
    ) 
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # 开始训练
    print(f"开始训练")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, return_tensors='pt')
    trainer = Seq2SeqTrainer(
        model=model,
        tokenizer=tokenizer,
        args=args,
        train_dataset=tra_ds,
        eval_dataset=val_ds,
        data_collator=data_collator,
    )
    trainer.train()
    trainer.save_model(os.path.join(ws, f"models/sft/{model_name}"))