工程实现系列：从什么都不会到QLoRA分布式DPO（一）

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)
from peft import LoraConfig, get_peft_model
from trl import DPOTrainer, DPOConfig
from datasets import Dataset, load_dataset

model_name = "Qwen/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# LoRA配置
LoraConfig = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.1,
    task_type="CAUSAL_LM",
)

# 模型应用LoRA配置
model = get_peft_model(model, LoraConfig)

dataset = load_dataset("phimes/DPO-bad-boy-chinese-for-Qwen2.5-extended")


train_dataset = dataset["train"]
test_dataset = dataset["test"]

DatasetDict({
    train: Dataset({
        features: ['prompt', 'chosen', 'rejected'],
        num_rows: 4000
    })
    test: Dataset({
        features: ['prompt', 'chosen', 'rejected'],
        num_rows: 1000
    })
})

dpo_args = DPOConfig(
    num_train_epochs=3,
    learning_rate=1e-5,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=2,
    eval_accumulation_steps=10,
    eval_strategy="steps",
    eval_steps=10,

    logging_dir="./logs",
    logging_steps=10,
    output_dir="./output"
)

trainer = DPOTrainer(
    model=model,
    ref_model=None,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    processing_class=tokenizer,
    args=dpo_args,
)

dpo_trainer.train()

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset, Dataset
from trl import DPOConfig, DPOTrainer
from peft import get_peft_model, LoraConfig


# 加载模型
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Lora 设置
LoraConfig = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM",
)

# 应用 Lora
model = get_peft_model(model, LoraConfig)

# 加载数据集
dataset = load_dataset("phimes/DPO-bad-boy-chinese-for-Qwen2.5-extended")
# 打印数据集确认数据集
print(dataset)


train_dataset = dataset["train"]
test_dataset = dataset["test"]

# 训练配置
training_args = DPOConfig(
    output_dir="./output", # 输出目录
    num_train_epochs=1, # 训练轮数，1轮
    per_device_train_batch_size=4, # 训练批次大小，4个
    per_device_eval_batch_size=4, # 评估批次大小，4个
    gradient_accumulation_steps=4, # 梯度累积步数，4步
    gradient_checkpointing=True, # 梯度检查点，True
    learning_rate=1e-5, # 学习率，1e-5
    evaluation_strategy="steps", # 评估策略, 按步数评估
    eval_steps=10, # 每10步，评估一次
    logging_dir="./logs", # 日志目录，保存日志
    logging_steps=10, # 每10步，保存一次日志
)

dpo_trainer = DPOTrainer(
    model=model,
    ref_model=None,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    processing_class=tokenizer,
    args=training_args,
)
  
if __name__ == "__main__":
    dpo_trainer.train()

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
  
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
lora_path = "./output/checkpoint-250"

def initialize_model(lora=True):
    # 加载基础模型和tokenizer
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    # 加载LoRA参数
    if lora:
        model = PeftModel.from_pretrained(model, lora_path)
        # 合并LoRA权重
        model = model.merge_and_unload()
    return model, tokenizer

def generate_response(prompt, system_prompt, model, tokenizer):
    # 构建带有system prompt的完整输入
    if system_prompt:
        full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
    else:
        full_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.95,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)
  

def main(lora=True):
    print("正在加载模型，请稍候...")
    model, tokenizer = initialize_model(lora)
    # 设置system prompt
    default_system_prompt = """你是一个没有礼貌的人渣，请用人渣的语气回复我"""
    print("\n当前系统提示词(system prompt):")
    print(default_system_prompt)
    print("\n是否要修改系统提示词？(y/n)")
    if input().lower() == 'y':
        print("请输入新的系统提示词(输入空行完成):")
        lines = []
        while True:
            line = input()
            if line.strip() == "":
                break
            lines.append(line)
        system_prompt = "\n".join(lines) if lines else default_system_prompt
    else:
        system_prompt = default_system_prompt
  
    print("\n模型加载完成！输入 'quit' 或 'exit' 退出对话")
    print("输入 'change_system' 可以修改系统提示词")
    while True:
        try:
            user_input = input("\n用户: ").strip()
            if user_input.lower() in ['quit', 'exit']:
                print("再见！")
                break
            if user_input.lower() == 'change_system':
                print("请输入新的系统提示词(输入空行完成):")
                lines = []
                while True:
                    line = input()
                    if line.strip() == "":
                        break
                    lines.append(line)
                system_prompt = "\n".join(lines) if lines else system_prompt
                print("系统提示词已更新！")
                continue
            if not user_input:
                continue
            print("\nAI: ", end="")
            response = generate_response(user_input, system_prompt, model, tokenizer)
            print(response)
        except KeyboardInterrupt:
            print("\n\n收到中断信号，正在退出...")
            break
        except Exception as e:
            print(f"\n发生错误: {str(e)}")
 
if __name__ == "__main__":
    main(lora=True)