import os
import pickle
from transformers import AutoModelForCausalLM
from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, TaskType


device = "cuda:0"

#1.创建lora微调基本的配置
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

 2.通过调用get_peft_model方法包装基础的Transformer模型

#通过调用get_peft_model方法包装基础的Transformer模型
model = AutoModelForCausalLM("/root/paddlejob/workspace/llama-2-7b-chat")
model = get_peft_model(model, peft_config)

下面是lora微调的模型结构，可以看到多了两个矩阵，一个降维一个升维 

3.训练

# optimizer and lr scheduler
'''len(train_dataloader) 是训练数据集中的批次数量，num_epochs 是训练过程中的迭代次数。因此，len(train_dataloader) * num_epochs 表示整个训练过程中的总迭代次数，即总共要遍历训练数据集的批次数'''
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=(len(train_dataloader) * num_epochs),)

#training and evaluation
model = model.to('cuda:0')
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for step, batch in enumerate(train_dataloader):
        batch = {k: v.to('cuda:0') for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        total_loss = total_loss + loss.detach().float()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    model.eval()
    eval_loss = 0
    eval_preds = []
   for step, batch in enumerate(eval_dataloader):
        batch = {k: v.to('cuda:0') for k, v in bacth.items()}
        with torch.nograd():
            outputs = model(**bacth)
        loss = outputs.loss
        eval_loss = eval_loss + loss.detach().float()
        eval_preds.extend(
            tokenizer.bacth_decode(torch.argmax(outputs.logits, -1), skip_special_tokens=True)
            )
    
    eval_epoch_loss = eval_loss / len(eval_dataloader)
    eval_ppl = torch.exp(eval_epoch_loss)
    train_epoch_loss = total_loss / len(train_dataloader)
    train_ppl = torch.exp(train_epoch_loss)
        

'''
在 Python 中，** 运算符用于将字典解包为关键字参数传递给函数或方法。在 PyTorch 中，model(**batch) 中的 **batch 将字典 batch 中的键值对作为关键字参数传递给模型的方法（通常是前向传播方法）。

具体来说，**batch 将字典 batch 中的每个键值对解包为一组关键字参数。例如，如果 batch 字典包含键值对 {'input_ids': tensor1, 'attention_mask': tensor2}，那么 model(**batch) 实际上就等价于 model(input_ids=tensor1, attention_mask=tensor2)。

这种方式可以方便地将字典中的数据传递给函数或方法，并且使代码更加简洁和易读。
'''

'''
loss.detach().float() 的作用是将计算图中的 loss 张量分离出来并转换为浮点数类型。

具体来说：

loss.detach() 会创建一个新的张量，其值与 loss 相同，但不再跟踪梯度信息。这样做是因为在训练过程中，我们通常只需要保存当前步骤的损失值，而不需要其相关的计算图和梯度信息。
.float() 将张量转换为浮点数类型。这是因为通常情况下，损失值是作为浮点数来计算和累加的。
所以，total_loss += loss.detach().float() 的作用就是将当前步骤的损失值添加到总损失值中，保证总损失值是一个浮点数。
'''

'''
在使用 PyTorch 进行梯度下降优化时，optimizer.zero_grad() 的作用是将模型参数的梯度归零，以便进行新一轮的梯度计算和更新。这是因为在 PyTorch 中，每次调用 .backward() 方法都会累积梯度，而不是覆盖之前的梯度。因此，在每次迭代更新参数之前，需要先将之前的梯度清零，以免影响当前迭代的梯度计算。

简而言之，optimizer.zero_grad() 用于初始化梯度，确保每次迭代都是基于当前 batch 的梯度计算和参数更新，而不会受到之前迭代的影响。
'''

'''

outputs.logits 是模型生成的原始输出，通常是一个三维张量，其中包含了模型对于每个词汇的得分（未经过 softmax 处理）。在语言模型中，这个张量的维度通常是 (batch_size, sequence_length, vocab_size)，其中 batch_size 表示批量大小，sequence_length 表示每个序列的长度，vocab_size 表示词汇表的大小。
在生成文本任务中，outputs.logits 的每个元素表示模型在当前位置生成每个词汇的得分。通常，需要对这些得分进行 softmax 处理以获得每个词汇的概率分布，然后根据概率分布进行采样或选择最高概率的词汇作为模型生成的下一个词。

torch.argmax 是 PyTorch 库中的一个函数，用于返回张量中指定维度上的最大值的索引。具体而言，对于一个输入张量，torch.argmax(input, dim=None, keepdim=False) 函数将返回指定维度 dim 上最大值的索引。如果不指定 dim，则默认返回整个张量中最大值的索引。
例如，对于一个形状为 (batch_size, seq_length, vocab_size) 的张量，torch.argmax(outputs.logits, -1) 将返回在 vocab_size 维度上每个位置上的最大值对应的索引，即得分最高的词的索引。

这行代码的作用是将模型输出的 logits （对应每个词的得分）经过 torch.argmax 函数找到得分最高的词的索引，然后使用 tokenizer 对这些索引进行解码，将索引转换为对应的词，并通过 skip_special_tokens=True 参数去除特殊标记（如 [CLS], [SEP] 等）。最终得到的是模型生成的文本内容。
'''

 4.模型保存

#save model
peft_model_id = f"{model_name_path}_{peft_config.peft_type}_{peft_config.peft.task_type}"
model.save_pretrained(peft_model_id)

5.模型训练的其余部分无需更改，当模型训练完成后，保存高效微调的模型权重部分以供模型推理 

#加载微调后的权重
from peft import PeftModel, PeftConfig

config = PeftConfig.from_pretrained(peft_model_id)
##加载基础模型
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
##加载peft模型
model = PeftModel.from_pretrained(model, peft_model_id)

##加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
tokennizer.pad_token = tokenizer.eos_token

6.加载微调后的权重文件，并进行推理 

#利用微调后的模型进行推理
##tokenizer编码
inputs = tokenizer(f'{text_column} : {dataset["test"][i]["Tweet text"]} Label : ', return_tensors="pt")

##模型推理
outputs = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs["attention_mask"],
    max_new_tokens=10,
    eos_token_id=3
)

##tokenizer解码
print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))