LoRA 方法介绍¶

# ---------------------------
# 1️⃣ 导入必要库
# ---------------------------
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
from peft import LoraConfig, get_peft_model
from datasets import load_dataset

# ---------------------------
# 2️⃣ 加载 tokenizer
# ---------------------------
tokenizer = AutoTokenizer.from_pretrained(...)

# ---------------------------
# 3️⃣ 加载并处理数据集
# ---------------------------
dataset = load_dataset(...)

# ---------------------------
# 4️⃣ 加载原模型（冻结大部分参数）
# ---------------------------
model = AutoModelForCausalLM.from_pretrained(...)

# ---------------------------
# 5️⃣ 配置 LoRA 微调
# ---------------------------
lora_config = LoraConfig(...)

model = get_peft_model(model, lora_config)

# ---------------------------
# 6️⃣ 数据整理器（batch）
# ---------------------------
data_collator = DataCollatorForSeq2Seq(...)

# ---------------------------
# 7️⃣ 配置训练参数
# ---------------------------
training_args = TrainingArguments(...)

# ---------------------------
# 8️⃣ Trainer
# ---------------------------
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    data_collator=data_collator
)

# ---------------------------
# 9️⃣ 开始训练
# ---------------------------
trainer.train()

# ---------------------------
# 🔟 保存 LoRA 权重
# ---------------------------
model.save_pretrained(...)

def load_tokenizer(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    return tokenizer

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")

example = {
    "Context": "Once upon a time, there was a brave knight",
    "Response": " who fought dragons to protect the kingdom."
}

# 将输入和输出拼接
input_text = example["Context"]
target_text = example["Response"]

input_ids = tokenizer(input_text, truncation=True, max_length=128).input_ids
target_ids = tokenizer(target_text, truncation=True, max_length=128).input_ids

# 构造 labels：只计算 target loss
labels = [-100]*len(input_ids) + target_ids

# attention mask
attention_mask = [1]*len(labels)

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct", trust_remote_code=True)

SYSTEM_PROMPT = "You are a helpful mental health assistant."

example = {
    "Context": "I feel anxious about my exams.",
    "Response": "It's normal to feel anxious. Try some breathing exercises..."
}

# 构造 messages
messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": example["Context"]},
    {"role": "assistant", "content": example["Response"]}
]

# 使用 apply_chat_template
ids = tokenizer.apply_chat_template(messages, tokenize=True)

# 假设返回两个部分：prompt_ids 和 full_ids
prompt_ids = tokenizer.apply_chat_template(messages[:-1], tokenize=True, add_generation_prompt=True)
full_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)
target_ids = full_ids[len(prompt_ids):]

input_ids = prompt_ids + target_ids
labels = [-100]*len(prompt_ids) + target_ids
attention_mask = [1]*len(input_ids)

history = [
    {"role": "user", "content": "Hello!"},
    {"role": "assistant", "content": "Hi, how can I help you today?"}
]

current_input = "I feel stressed about work."
current_response = "It is common to feel stressed. Let's talk through it."

messages = [{"role": "system", "content": SYSTEM_PROMPT}] + history + \
           [{"role": "user", "content": current_input}, {"role": "assistant", "content": current_response}]

# 处理同 apply_chat_template
ids = tokenizer.apply_chat_template(messages, tokenize=True)
prompt_ids = tokenizer.apply_chat_template(messages[:-1], tokenize=True)
full_ids = tokenizer.apply_chat_template(messages, tokenize=True)
target_ids = full_ids[len(prompt_ids):]

input_ids = prompt_ids + target_ids
labels = [-100]*len(prompt_ids) + target_ids
attention_mask = [1]*len(input_ids)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B-Instruct",
    device_map="auto",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B-Instruct",
    device_map="auto",
    quantization_config=quant_config,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)

def apply_lora(model):
    lora_config = LoraConfig(
        task_type="CAUSAL_LM",
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"]
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    return model

data_collator = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,
    model=model,
    padding=True,
    label_pad_token_id=-100,
    return_tensors="pt"
)

training_args = TrainingArguments(
    output_dir="./qwen2.5_0.5B_qlora_4060_val_1e-4",
    overwrite_output_dir=True,
    do_train=True,
    do_eval=True,
    eval_strategy="steps",
    eval_steps=100,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    learning_rate=2e-4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    max_grad_norm=0.3,
    logging_strategy="steps",
    logging_steps=10,
    logging_dir="./qwen2.5_0.5B_qlora_4060_val_1e-4/logs",
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    save_safetensors=True,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    bf16=torch.cuda.is_available(),
    fp16=not torch.cuda.is_bf16_supported() and torch.cuda.is_available(),
    optim="paged_adamw_8bit",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    group_by_length=True,
    report_to=["tensorboard"],
    remove_unused_columns=False
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_eval,
    data_collator=data_collator
)

# 从头开始训练
trainer.train()
# 继续 ./qwen2.5_0.5B_qlora_4060_val_1e-4 下保存的 checkpoint-1100 节点
trainer.train(resume_from_checkpoint="./qwen2.5_0.5B_qlora_4060_val_1e-4/checkpoint-1100")

trainer.model.save_pretrained("./qwen2.5_0.5B_qlora_4060_val_1e-4") # 保存 adapter 权重和 config  
tokenizer.save_pretrained("./qwen2.5_0.5B_qlora_4060_val_1e-4") # 保存 tokenizer 配置

"""
单卡 4060 QLoRA 微调脚本（带验证集），训练参数集中在配置区
适用于心理咨询对话数据集
"""

import importlib.util

import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    Trainer,
    TrainingArguments,
    DataCollatorForSeq2Seq,
    set_seed,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# --------------------------- 配置区域 ---------------------------

# ---------------- 基础模型与数据集 ----------------
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"   # 基础预训练模型名称
DATASET_NAME = "Amod/mental_health_counseling_conversations"  # 数据集名称
OUTPUT_BASE_DIR = "./qwen2.5_0.5B_qlora_4060_val"      # LoRA adapter 和 tokenizer 保存路径
SEED = 42                                 # 随机种子，保证训练可复现

TASK_TYPE = "CAUSAL_LM"                   #  LoRA 任务类型，需要与基础与训练模型类型保持一直

# ---------------- LoRA 参数（影响微调效果） ----------------
LORA_R = 8                # LoRA 秩 r，越大模型可训练容量越高，但显存占用也越大
LORA_ALPHA = 32           # LoRA 缩放系数 alpha，通常 alpha/r ≈ 实际缩放
LORA_DROPOUT = 0.05       # LoRA 分支 dropout，用于防止小数据过拟合
TARGET_MODULES = [        # 注入 LoRA 的模型线性层
    "q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"
]

# ---------------- 训练参数 ----------------
EVAL_STEPS = 100          # 每多少训练步进行一次验证
NUM_EPOCHS = 3            # 数据集遍历轮数
TRAIN_BATCH_SIZE = 1      # 每张 GPU 训练 batch 大小
EVAL_BATCH_SIZE = 1       # 每张 GPU 验证 batch 大小
GRAD_ACCUM_STEPS = 8      # 梯度累积步数（多步梯度后再更新一次参数）
LEARNING_RATE = 8e-4      # 初始学习率
WARMUP_RATIO = 0.03       # 学习率预热比例（训练前期慢慢升到学习率）
LR_SCHEDULER_TYPE = "cosine"  # 学习率调度策略
MAX_GRAD_NORM = 0.3       # 梯度裁剪阈值，避免梯度爆炸
LOGGING_STEPS = 10        # 每隔多少个 step 打印一次训练日志
SAVE_STEPS = 100          # 每隔多少个 step 保存一次 checkpoint
SAVE_TOTAL_LIMIT = 2      # 最多保留多少个 checkpoint，旧的会被清理
OPTIMIZER = "paged_adamw_8bit" # 优化器

# ---------------- 根据学习率设置输出路径 ----------------
OUTPUT_DIR = f"{OUTPUT_BASE_DIR}_{'{:.0e}'.format(LEARNING_RATE).replace('e-0', 'e-')}"

# ---------------- 训练日志与可视化 ----------------
USE_TENSORBOARD = True    # 使用 TensorBoard 查看训练 loss、eval_loss、learning_rate 等曲线
LOGGING_DIR = f"{OUTPUT_DIR}/logs"

# ---------------- 验证集参数 ----------------
EVAL_WEIGHT = 0.1         # 验证集样本占比

# ---------------- 文本处理参数 ----------------
MAX_SEQ_LENGTH = 768      # 输入文本最大长度（超过长度截断）

# ---------------- 断点续训 ----------------
RESUME_FROM_CHECKPOINT = None   # 添加保存 checkpoint 的文件夹路径

# ---------------- 系统提示 ----------------
SYSTEM_PROMPT = (
    "你是一名富有同理心的心理健康咨询助手。"
    "请用英语回答，语气温暖、能体现倾听，并给出可执行的下一步建议。"
    "不要做诊断。若问题严重、持续存在或涉及安全风险，应鼓励对方寻求合格专业人士的帮助。"
)

# --------------------------- 工具函数 ---------------------------

def set_seed_everywhere(seed):
    set_seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def load_tokenizer(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    return tokenizer

def build_messages(context, response=None):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": context.strip()},
    ]
    if response is not None:
        messages.append({"role": "assistant", "content": response.strip()})
    return messages

def tokenize_example(example, tokenizer, max_seq_length=MAX_SEQ_LENGTH):
    prompt_ids = tokenizer.apply_chat_template(build_messages(example["Context"]), tokenize=True, add_generation_prompt=True)
    full_ids = tokenizer.apply_chat_template(build_messages(example["Context"], example["Response"]), tokenize=True, add_generation_prompt=False)
    if full_ids[:len(prompt_ids)] != prompt_ids:
        raise ValueError("chat template 前缀不匹配")
    target_ids = full_ids[len(prompt_ids):]
    if len(prompt_ids) + len(target_ids) > max_seq_length:
        min_target_tokens = min(256, len(target_ids), max_seq_length-1)
        max_prompt_tokens = max_seq_length - min_target_tokens
        prompt_ids = prompt_ids[-max_prompt_tokens:]
        target_ids = target_ids[:max_seq_length - len(prompt_ids)]
    input_ids = prompt_ids + target_ids
    labels = [-100]*len(prompt_ids) + target_ids
    return {"input_ids": input_ids, "attention_mask": [1]*len(input_ids), "labels": labels}

def prepare_dataset(dataset_name):
    ds = load_dataset(dataset_name, split="train")
    ds = ds.filter(lambda x: bool(x["Context"].strip()) and bool(x["Response"].strip()))
    ds = ds.map(lambda x: {"Context": x["Context"].strip(), "Response": x["Response"].strip()})
    return ds

def apply_lora(model):
    lora_config = LoraConfig(
        task_type=TASK_TYPE,
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        target_modules=TARGET_MODULES
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    return model

def build_quantization_config():
    return BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    )

# --------------------------- 主流程 ---------------------------

def main():
    set_seed_everywhere(SEED)
    if not torch.cuda.is_available():
        raise RuntimeError("QLoRA 微调需要 GPU")

    # Tokenizer
    tokenizer = load_tokenizer(MODEL_NAME)

    # 模型
    quant_config = build_quantization_config()
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        quantization_config=quant_config,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    model.config.use_cache = False
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    model = apply_lora(model)

    # 数据集 + 切分验证集
    dataset = prepare_dataset(DATASET_NAME)
    split = dataset.train_test_split(test_size=0.1, seed=SEED, shuffle=True)
    train_dataset = split["train"]
    eval_dataset = split["test"]

    tokenized_train = train_dataset.map(lambda x: tokenize_example(x, tokenizer), remove_columns=dataset.column_names)
    tokenized_eval = eval_dataset.map(lambda x: tokenize_example(x, tokenizer), remove_columns=dataset.column_names)
    print(f"总数据集：{len(dataset)}， 训练数据集：{len(tokenized_train)}， 验证数据集：{len(tokenized_eval)}")

    # DataCollator
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, padding=True, label_pad_token_id=-100, return_tensors="pt")

    # Trainer
    if USE_TENSORBOARD and importlib.util.find_spec("tensorboard") is None:
        raise RuntimeError("USE_TENSORBOARD=True 需要先安装 tensorboard：python -m pip install tensorboard")

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        overwrite_output_dir=True,
        do_train=True,
        do_eval=True,
        eval_strategy="steps",
        eval_steps=EVAL_STEPS,
        per_device_train_batch_size=TRAIN_BATCH_SIZE,
        per_device_eval_batch_size=EVAL_BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM_STEPS,
        num_train_epochs=NUM_EPOCHS,
        learning_rate=LEARNING_RATE,
        warmup_ratio=WARMUP_RATIO,
        lr_scheduler_type=LR_SCHEDULER_TYPE,
        max_grad_norm=MAX_GRAD_NORM,
        logging_strategy="steps",
        logging_steps=LOGGING_STEPS,
        logging_dir=LOGGING_DIR,
        save_strategy="steps",
        save_steps=SAVE_STEPS,
        save_total_limit=SAVE_TOTAL_LIMIT,
        save_safetensors=True,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        bf16=torch.cuda.is_available(),
        fp16=not torch.cuda.is_bf16_supported() and torch.cuda.is_available(),
        optim=OPTIMIZER,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        group_by_length=True,
        report_to=["tensorboard"] if USE_TENSORBOARD else "none",
        remove_unused_columns=False
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train,
        eval_dataset=tokenized_eval,
        data_collator=data_collator
    )

    print("开始训练...")
    trainer.train(resume_from_checkpoint=RESUME_FROM_CHECKPOINT)
    print("训练完成")
    print(f"最佳 checkpoint: {trainer.state.best_model_checkpoint}")
    print(f"最佳 eval_loss: {trainer.state.best_metric}")

    # 保存
    trainer.model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    print(f"模型和 LoRA adapter 已保存到 {OUTPUT_DIR}")

if __name__ == "__main__":
    main()

"""
通用 LoRA / QLoRA 微调效果评估脚本
功能：
1. 加载微调后的模型和 tokenizer
2. 加载验证/测试数据集
3. 使用模型生成回答
4. 计算生成回答与参考答案的相似度（cosine similarity）
"""

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from datasets import load_dataset
from sentence_transformers import SentenceTransformer, util
from peft import PeftModel, prepare_model_for_kbit_training
from tqdm.auto import tqdm

# --------------------------- 配置 ---------------------------
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"

# 微调后的模型路径
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_1e-4"
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_2e-4"
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_3e-4"
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_4e-4"
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_5e-4"
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_6e-4"
# MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_7e-4"
MODEL_DIR = "./qwen2.5_0.5B_qlora_4060_val_8e-4"

DATASET_NAME = "Amod/mental_health_counseling_conversations"  # 可替换为其他本地/Hub数据集
SEED = 42                                 # 随机种子，保持与微调脚本一致
MAX_SEQ_LENGTH = 768
MAX_NEW_TOKENS = 512
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

SYSTEM_PROMPT = (
    "你是一名富有同理心的心理健康咨询助手。"
    "请用英语回答，语气温暖、能体现倾听，并给出可执行的下一步建议。"
    "不要做诊断。若问题严重、持续存在或涉及安全风险，应鼓励对方寻求合格专业人士的帮助。"
)

# --------------------------- 加载模型和 tokenizer ---------------------------

def load_tokenizer(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    return tokenizer

def build_quantization_config():
    return BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    )
# Tokenizer
tokenizer = load_tokenizer(MODEL_NAME)

# 模型
quant_config = build_quantization_config()
base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    quantization_config=quant_config,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)

base_model.config.use_cache = True

# model = base_model
model = PeftModel.from_pretrained(base_model, MODEL_DIR)

model.to(DEVICE)
model.eval()

# --------------------------- 加载验证数据集 ---------------------------
dataset = load_dataset(DATASET_NAME, split="train")
dataset = dataset.filter(lambda x: bool(x["Context"].strip()) and bool(x["Response"].strip()))
dataset = dataset.map(lambda x: {"Context": x["Context"].strip(), "Response": x["Response"].strip()})

split = dataset.train_test_split(test_size=0.1, seed=SEED, shuffle=True)
dataset = split["test"]

# --------------------------- 嵌入模型 ---------------------------
# 用 sentence-transformers 计算文本相似度
embedder = SentenceTransformer('all-mpnet-base-v2', device=DEVICE)

# --------------------------- 生成函数 ---------------------------
def generate_answer_with_template(msg: list, max_new_tokens: int = MAX_NEW_TOKENS) -> str:
    inputs = tokenizer.apply_chat_template(
        msg,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        truncation=True,
        max_length=MAX_SEQ_LENGTH,
    ).to(DEVICE)
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    # 去掉 prompt 部分
    prompt_len = inputs["input_ids"].shape[-1]
    generated_ids = output_ids[0][prompt_len:]
    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

# --------------------------- 测试与相似度计算 ---------------------------
results = []
similarity_total = 0.0
total_samples = len(dataset)
print(f"开始评估：{total_samples} 条样本，模型：{MODEL_DIR if model is not base_model else MODEL_NAME}")

progress_bar = tqdm(dataset, total=total_samples, desc="Evaluating", unit="sample")
for idx, example in enumerate(progress_bar, start=1):
    # 构造 messages
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": example['Context'].strip()},
    ]
    reference = example['Response']

    # 生成模型回答
    generated = generate_answer_with_template(messages)

    # 计算文本嵌入
    emb_generated = embedder.encode(generated, convert_to_tensor=True)
    emb_reference = embedder.encode(reference, convert_to_tensor=True)

    # 余弦相似度
    similarity = util.cos_sim(emb_generated, emb_reference).item()

    results.append({
        "Context": example['Context'],
        "Reference": reference,
        "Generated": generated,
        "Similarity": similarity
    })

    similarity_total += similarity
    running_avg = similarity_total / idx
    progress_bar.set_postfix(
        current=f"{similarity:.4f}",
        avg=f"{running_avg:.4f}",
    )

# --------------------------- 输出结果 ---------------------------
# for r in results:
#     print("------")
#     print("Context:", r["Context"])
#     print("Reference:", r["Reference"])
#     print("Generated:", r["Generated"])
#     print(f"Similarity: {r['Similarity']:.4f}")

# 平均相似度
avg_sim = sum(r["Similarity"] for r in results)/len(results)
print(f"\nAverage Similarity: {avg_sim:.4f}")

trainer.model.save_pretrained("./lora_adapter")

model = PeftModel.from_pretrained(base_model, adapter_path)
merged_model = model.merge_and_unload()

merged_model.save_pretrained("./merged_model", safe_serialization=True)
tokenizer.save_pretrained("./merged_model")

"""
多 LoRA adapter 加载与切换演示
"""
from pathlib import Path

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


# --------------------------- 配置 ---------------------------
BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"

# 这里分别定义三个 adapter,用来演示多 adapter 下如何加载和切换 adapter
ADAPTERS = {
    "one": Path("./qwen2.5_0.5B_qlora_4060_val_1e-4"),
    "two": Path("./qwen2.5_0.5B_qlora_4060_val_4e-4"),
    "three": Path("./qwen2.5_0.5B_qlora_4060_val_7e-4"),
}

# 为每个 adapter 定义 System Prompt
SYSTEM_PROMPTS = {
    "one": "你是一个专业、温和、谨慎的心理咨询助手。",
    "two": "你是一个专业、温和、谨慎的心理咨询助手。",
    "three": "你是一个专业、温和、谨慎的心理咨询助手。",
}


def load_multi_adapter_model():
    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL,
        trust_remote_code=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
    )
    base_model.config.use_cache = True

    first_name = next(iter(ADAPTERS))
    first_path = ADAPTERS[first_name]

    model = PeftModel.from_pretrained(
        base_model,
        first_path,
        adapter_name=first_name,
    )

    for adapter_name, adapter_path in ADAPTERS.items():
        if adapter_name == first_name:
            continue

        model.load_adapter(
            adapter_path,
            adapter_name=adapter_name,
        )

    model.eval()
    return tokenizer, model


def generate(
    tokenizer,
    model,
    adapter_name: str,
    query: str,
    max_new_tokens: int = 512,
) -> str:
    if adapter_name not in ADAPTERS:
        raise ValueError(f"未知 adapter_name: {adapter_name}")

    model.set_adapter(adapter_name)

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPTS.get(adapter_name, "你是一个专业助手。"),
        },
        {
            "role": "user",
            "content": query.strip(),
        },
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
    # 去掉 prompt 部分
    prompt_len = inputs["input_ids"].shape[-1]
    generated_ids = output_ids[0][prompt_len:]
    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()


def main():

    tokenizer, model = load_multi_adapter_model()

    examples = [
        ("one", "我最近压力很大，晚上睡不着怎么办？"),
        ("two", "我最近压力很大，晚上睡不着怎么办？"),
        ("three", "我最近压力很大，晚上睡不着怎么办？"),
    ]

    for adapter_name, query in examples:
        answer = generate(tokenizer, model, adapter_name, query)
        print("=" * 80)
        print(f"Adapter: {adapter_name}")
        print(f"User: {query}")
        print(f"Assistant: {answer}")

if __name__ == "__main__":
    main()