import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 1. 基础配置
model_path = "./model"  # 你之前下载的本地模型路径
output_dir = "./lora_results" # 微调后 LoRA 权重的保存位置

print("1. 正在加载模型与分词器...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
# 微调时通常需要 pad_token 来对齐数据长度。如果模型没有，指定 eos_token 为 pad_token
# 注意：微调时通常需要 pad_token。如果模型没有，我们指定 eos_token 为 pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 加载模型，这里我们为了节省内存，可以加载为 16 位浮点数 (bfloat16)
# device_map="auto" 会自动将模型分配到可用的计算设备（如 Mac 的 mps，或 Nvidia 的 cuda）
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto" # 自动分配到 mps/cuda/cpu
)

# 加载刚刚创建的 json 数据
dataset = load_dataset("json", data_files="dataset.json", split="train")

# 定义一个将 instruction 和 output 拼接为对话模板的函数
def format_prompt(example):
    # 这里模拟模型的对话格式
    # 构建对话上下文，末尾加上 eos_token 告诉模型回答结束
    text = f"User: {example['instruction']}\nAssistant: {example['output']}{tokenizer.eos_token}"
    return {"text": text}

# 映射数据集
formatted_dataset = dataset.map(format_prompt)

# 【关键学习点】这就是 LoRA 的魔法所在！我们不训练模型原本的上百亿参数，而是插入几个小矩阵。
lora_config = LoraConfig(
    r=8,               # Rank (秩)：决定了新增矩阵的大小。值越大，模型能学到的细节越多，但内存消耗也越大。通常 8, 16, 32 都可以。
    lora_alpha=16,     # 缩放系数：控制 LoRA 权重对原模型的影响力度，通常设为 r 的 2 倍。
    target_modules=["q_proj", "v_proj"], # 要在哪些层应用 LoRA？通常在注意力机制的 Query 和 Value 矩阵上效果最好。
    lora_dropout=0.05, # 防过拟合：随机丢弃 5% 的神经元。
    bias="none",
    task_type="CAUSAL_LM" # 任务类型：因果语言建模（即文本生成）
)

# 将普通模型包装成支持 LoRA 训练的模型
model = get_peft_model(model, lora_config)
# 可以通过 model.print_trainable_parameters() 查看，你会发现可训练参数通常不到 1%

training_args = SFTConfig(
    output_dir=output_dir,
    per_device_train_batch_size=2,  # 每次处理的数据量，如果内存/显存爆了，就把这个调成 1
    gradient_accumulation_steps=2,  # 梯度累加，变相增加 Batch Size，节省内存
    learning_rate=2e-4,  # 学习率
    num_train_epochs=3,  # 训练轮数：把数据集看几遍
    logging_steps=1,  # 每隔多少步打印一次日志
    save_strategy="epoch",  # 保存策略
    optim="adamw_torch",
    dataset_text_field="text", # 指定数据集中作为训练文本的列名
)


# SFTTrainer 会帮我们自动处理文本的 Tokenize 和 Padding，非常省心
trainer = SFTTrainer(
    model=model,
    train_dataset=formatted_dataset,
#     max_seq_length=512,
    args=training_args,
)

# 开始训练！
print("开始注入领域知识...")
trainer.train()

# 保存模型和分词器（注意：这里只保存了极小的 LoRA 权重，即“便利贴”）
trainer.model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
print(f"训练完成！LoRA 权重已保存至 {output_dir}")

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# 1. 定义两个路径
base_model_path = "./model"         # 你最初下载的、几十亿参数的基础大模型路径
lora_weights_path = "./lora_results" # 你刚才微调保存的 "便利贴" (补丁) 路径

print("1. 正在加载基础模型和分词器 (搬出百科全书)...")
tokenizer = AutoTokenizer.from_pretrained(base_model_path)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto" # 自动调用 Mac 的 mps 加速
)

print("2. 正在融合微调权重 (贴上便利贴)...")
# 【核心魔法】将基础模型和 LoRA 补丁组装在一起
model = PeftModel.from_pretrained(base_model, lora_weights_path)
print("模型融合完毕！\n")

rint("=" * 50)
print("欢迎测试微调后的专属助手！输入 'quit' 结束。")
print("=" * 50)

# 3. 开始连续对话测试
while True:
    user_input = input("\nUser: ").strip()
    if user_input.lower() in ['quit', 'exit']:
        break
    if not user_input:
        continue

    # 【避坑指南】微调后的推理，必须严格使用训练时的对话格式！
    # 回顾我们刚才在 SFTTrainer 里写的格式：f"User: {instruction}\nAssistant: "
    prompt = f"User: {user_input}\nAssistant: "

    # 编码输入
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    input_length = inputs.input_ids.shape[1]

    # 模型推理生成
    outputs = model.generate(
        inputs.input_ids,
        attention_mask=inputs.attention_mask,
        max_new_tokens=256,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.3  # 降低温度，让模型回答更严谨
    )

    # 切片截取新生成的 Token
    new_tokens = outputs[0][input_length:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)

    # 过滤可能包含的思维链标签（针对特定推理模型）
    if "</think>" in response:
        clean_response = response.split("</think>")[-1].strip()
    else:
        clean_response = response.strip()
    # -----------------------

    print(f"\nAssistant: {clean_response}")

base_model_path = "./model/" # 原始底座，如 Qwen, Llama
sft_lora_path = "./lora_results/" # 你之前的 SFT 结果

# 1. 加载底座
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path, 
    torch_dtype=torch.float16, 
    device_map="auto"
)

# 2. 挂载 SFT LoRA
model = PeftModel.from_pretrained(base_model, sft_lora_path)

# 3. 合并并卸载 (核心步骤)
# 这会将 LoRA 的权重永久性地加到底座权重中
merged_model = model.merge_and_unload()

# 4. 保存为新的模型
merged_model.save_pretrained("./model_merged_for_rm")
tokenizer = AutoTokenizer.from_pretrained(base_model_path)
tokenizer.save_pretrained("./model_merged_for_rm")