from transformers import (AutoTokenizer,
                          AutoModelForCausalLM, #因果模型(自回归模型)
                          DataCollatorForLanguageModeling,#数据整理器, 用于将分词后的输入数据整理成模型训练所需要的格式
                          Trainer,#训练器
                          TrainingArguments
                          )
from datasets import load_dataset

#设置预训练模型的路径
model_name="/root/autodl-tmp/code/model"
#设置微调后模型保存的路径
finetuned_model_path="/root/autodl-tmp/code/fined"

#加载模型
# ---------- 第 1 步：配置 4bit 量化 ----------
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    local_files_only = True, #仅从本地文件加载模型
)

#加载预训练模型的tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
def preprocess_function(examples):
    """
    标准的 SFT 数据处理流程
    - 只训练 assistant 的回答部分
    - 正确处理截断和 EOS token
    - 保持 chat template 格式
    """
    max_length = 1024

    input_ids_list = []
    labels_list = []

    for question, cot, response in zip(
        examples["Question"],
        examples["Complex_CoT"],
        examples["Response"]
    ):
        # 构建对话消息
        messages = [
            {
                "role": "system",
                "content": "您是一位在临床推理、诊断和治疗计划方面具有高级知识的医学专家。请回答以下医学问题。"
            },
            {
                "role": "user",
                "content": question
            },
            {
                "role": "assistant",
                "content": f"<think>\n{cot}\n</think>\n{response}"
            }
        ]

        # ============ Step 1: 生成完整对话 ============
        full_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False
        )

        # ============ Step 2: 生成 prompt 部分（不含 assistant 回复）============
        prompt_messages = messages[:-1]  # 去掉最后的 assistant
        prompt_text = tokenizer.apply_chat_template(
            prompt_messages,
            tokenize=False,
            add_generation_prompt=True  # 添加 assistant 开始标记
        )

        # ============ Step 3: Tokenize ============
        full_ids = tokenizer.encode(full_text, add_special_tokens=False,
                                    max_length=max_length,padding="max_length", truncation=True)
        prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False,
                                      max_length=max_length,padding="max_length", truncation=True)

        # ============ Step 4: 处理截断（保留 EOS）============
        # apply_chat_template 已经自动添加了 EOS，所以只需要处理截断
        if len(full_ids) > max_length:
            # 截断到 max_length-1，然后手动添加 EOS
            full_ids = full_ids[:max_length - 1]
            full_ids.append(tokenizer.eos_token_id)
        # 如果没有超长，full_ids 已经包含 EOS，不需要额外处理

        # ============ Step 5: 构建 labels（关键！）============
        # prompt 部分设为 -100（不计算 loss）
        # assistant 部分保留原值（计算 loss）
        prompt_len = len(prompt_ids)
        labels = (
            [-100] * prompt_len +           # prompt 部分：不训练
            full_ids[prompt_len:]            # assistant 部分：训练（包含 EOS）
        )

        # 确保长度一致
        labels = labels[:len(full_ids)]

        input_ids_list.append(full_ids)
        labels_list.append(labels)

    return {
        "input_ids": input_ids_list,
        "attention_mask": [[1] * len(ids) for ids in input_ids_list],
        "labels": labels_list
    }

dataset = load_dataset(
    "json",
    data_files={
        "train": "/root/autodl-tmp/code/data/train.json",
        "valid": "/root/autodl-tmp/code/data/validation.json"
    }
)
dataset = dataset.map(preprocess_function, batched = True)

#设置微调模型参数
from peft import LoraConfig,get_peft_model,TaskType
#Lora配置
lora_config=LoraConfig(
    r=8,#低秩矩阵的秩,根据情况调整
    lora_alpha=16,#用于缩放Lora权重更新的超参数
    target_modules=['self_attn.q_proj','self_attn.k_proj','self_attn.v_proj','self_attn.o_proj'],
    lora_dropout=0.05,#dropout
    bias='none',#微调偏置
    task_type=TaskType.CAUSAL_LM
)
#获取应用了Lora后的微调模型
model = get_peft_model(model,lora_config)
model.config.use_cache = False
#实际上直接计算所有可训练参数就行。
trainable_params = 0
all_params = 0
for _, param in model.named_parameters():
    num_params = param.numel()
    all_params += num_params
    if param.requires_grad:
        trainable_params += num_params
print(f"可训练参数量: {trainable_params}")
print(f"总参数量: {all_params}")
print(f"可训练参数占比: {100 * trainable_params / all_params:.2f}%")

model.enable_input_require_grads() # 开启梯度检查点时，要执行该方法
#初始化可视化功能
#训练参数配置
training_args = TrainingArguments(
    output_dir=finetuned_model_path,#输出目录，用于保存模型和日志
    label_names=["labels"],  # 显式指定标签名称，确保 Trainer 正确识别 labels
    overwrite_output_dir=True,#覆盖输出目录的现有内容
    num_train_epochs=20,#训练的总轮次
    per_device_train_batch_size=6,#每个训练设备上的批次大小
    per_device_eval_batch_size=6,#每个评估设备上的批次大小
    save_steps=1,#每隔多少步保存一次模型
    save_total_limit=5,#保存的模型文件数量限制
    prediction_loss_only=True,#在评估期间只返回损失
    report_to="tensorboard",  # 关键：启用 tensorboard 上报
    logging_dir="./logs",
    save_strategy="epoch",
    #梯度累积步数
    #通过将多个小批量的梯度累积起来，然后一次性更新模型参数，从而在有限的内存资源下实现相当于更大批量的训练效果。
    #当累积的步数达到预设的 gradient_accumulation_steps 时，使用累积的梯度更新模型参数，并清零累积的梯度。
    #节省内存
    #模拟大批量训练
    #提高训练稳定性
    gradient_accumulation_steps=16, #梯度累积步数
    #梯度检查点，用于减少深度学习模型训练过程中内存占用
    #反向传播需要存储前向传播中每一层的中间激活值（activations），以便计算梯度。
    # 对于非常深的模型或大批量数据，这些激活值会占用大量内存
    #Gradient Checkpointing 的核心思想是只存储部分激活值，并在反向传播时重新计算其他激活值，从而减少内存占用。
    gradient_checkpointing=True,#启用梯度检查点
)

#创建模型数据整理器
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False #准备数据以进行标准语言模型任务
)

#创建训练器
trainer = Trainer(
    model = model,#要训练的模型
    args = training_args,#训练的参数
    data_collator=data_collator, # 用于准备训练数据
    train_dataset=dataset["train"],#训练的数据集
    eval_dataset=dataset["valid"] #评估的数据集
)

#开始训练模型
trainer.train()