构建AI智能体：大模型的指令微调与人类对齐：从知识渊博到善解人意

from datasets import Dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    TrainingArguments,
    BitsAndBytesConfig,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model
import torch
import os
from modelscope import snapshot_download

# 设置设备
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"使用设备: {device}")

# 1. 加载超小模型 - Qwen1.5-0.5B-Chat
model_name = "qwen/Qwen1.5-0.5B-Chat"
cache_dir = "D:\\modelscope\\hub"
print(f"正在加载模型: {model_name}")
local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True, cache_dir=cache_dir)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 根据设备选择配置
if device == "cuda":
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        local_model_path,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        cache_dir=cache_dir
    )
else:
    model = AutoModelForCausalLM.from_pretrained(
        local_model_path,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.float32,
        cache_dir=cache_dir
    )

print("模型加载完成!")

# 2. 配置优化的LoRA适配器
lora_config = LoraConfig(
    r=16,           # 适当增加秩以提高表达能力
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
print(f"可训练参数: {trainable_params:,} (仅占 {trainable_params/total_params*100:.2f}%)")

# 3. 准备优化的指令数据集
def prepare_instruction_dataset():
    base_dataset = [
        {
            "instruction": "翻译",
            "input": "Hello world → 中文",
            "output": "你好世界"
        },
        {
            "instruction": "情感分析", 
            "input": "太棒了！",
            "output": "积极"
        },
        {
            "instruction": "情感分析",
            "input": "真糟糕。", 
            "output": "消极"
        },
        {
            "instruction": "总结",
            "input": "今天天气晴朗，温度25度，适合出游。",
            "output": "天气晴好，适宜出行"
        },
        {
            "instruction": "写邮件",
            "input": "主题：会议取消",
            "output": "尊敬的与会者：\n\n原定会议已取消，具体安排另行通知。\n\n谢谢！"
        },
        {
            "instruction": "编程",
            "input": "Python Hello World",
            "output": "print('Hello World')"
        }
    ]
    
    # 扩展数据集 - 更多样化的指令
    expanded_dataset = base_dataset.copy()
    additional_examples = [
        {"instruction": "请翻译", "input": "Good morning → 中文", "output": "早上好"},
        {"instruction": "将以下英文翻译成中文", "input": "I love you", "output": "我爱你"},
        {"instruction": "翻译成中文", "input": "Thank you", "output": "谢谢"},
        {"instruction": "分析情感", "input": "我中奖了！", "output": "积极"},
        {"instruction": "这段话的情感是？", "input": "工作丢了...", "output": "消极"},
        {"instruction": "请总结", "input": "这款手机配备最新处理器，电池续航时间长，拍照效果很好。", "output": "手机性能强，续航久，拍照好"},
        {"instruction": "简明概括", "input": "今天会议讨论了新产品开发计划，决定下个月开始实施。", "output": "会议决定下月启动新产品开发"},
        {"instruction": "写邮件", "input": "主题：项目完成", "output": "尊敬的团队成员：\n\n项目已顺利完成，感谢大家的辛勤工作！\n\n祝好！"},
        {"instruction": "写代码", "input": "Python 计算两数之和", "output": "def add(a, b):\n    return a + b"},
        {"instruction": "回答问题", "input": "中国的首都是？", "output": "北京"},
        {"instruction": "解释概念", "input": "什么是人工智能？", "output": "人工智能是让机器模拟人类智能的技术。"}
    ]
    
    expanded_dataset.extend(additional_examples)
    return expanded_dataset

instruction_dataset = prepare_instruction_dataset()
print(f"总训练样本数: {len(instruction_dataset)}")

# 4. 优化的数据格式化函数
def format_instruction(sample):
    """使用ChatML格式，更适合Qwen模型"""
    if sample.get('input'):
        conversation = [
            {"role": "user", "content": f"{sample['instruction']}：{sample['input']}"},
            {"role": "assistant", "content": sample['output']}
        ]
    else:
        conversation = [
            {"role": "user", "content": sample['instruction']},
            {"role": "assistant", "content": sample['output']}
        ]
    
    # 使用tokenizer.apply_chat_template
    text = tokenizer.apply_chat_template(
        conversation, 
        tokenize=False, 
        add_generation_prompt=False
    )
    return text

# 准备训练数据
formatted_texts = [format_instruction(x) for x in instruction_dataset]

def tokenize_function(examples):
    # 启用 padding 和 truncation，确保数据长度一致
    tokenized = tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=512,
        return_tensors="pt"  # 返回 PyTorch 张量
    )
    
    # 对于因果语言建模，labels 就是 input_ids
    tokenized["labels"] = tokenized["input_ids"].clone()
    return tokenized

dataset = Dataset.from_dict({"text": formatted_texts})
tokenized_dataset = dataset.map(tokenize_function, batched=True)

print("数据预处理完成!")
print(f"样本示例: {formatted_texts[0][:100]}...")

# 5. 优化的训练配置
training_args = TrainingArguments(
    output_dir="./qwen1.5-0.5b-instruct-optimized",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=1e-4,  # 提高学习率
    num_train_epochs=15,  # 增加训练轮数
    logging_steps=5,
    save_steps=100,
    fp16=torch.cuda.is_available(),
    remove_unused_columns=True,
    dataloader_pin_memory=False,
    warmup_steps=10,  # 添加warmup
    evaluation_strategy="no",
    save_total_limit=2,
    report_to=None,  # 禁用wandb等记录
)

# 使用标准Trainer
from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    data_collator=DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,  # 因果语言建模
        pad_to_multiple_of=8
    )
)

# 6. 开始训练
print("开始优化训练...")
trainer.train()

# 保存完整模型（包括适配器）
trainer.save_model()
print("训练完成！模型已保存")

# 7. 优化的测试函数
def test_optimized_model(instruction, input_text=""):
    # 使用ChatML格式
    if input_text:
        messages = [
            {"role": "user", "content": f"{instruction}：{input_text}"}
        ]
    else:
        messages = [
            {"role": "user", "content": instruction}
        ]
    
    # 应用模板
    prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=True,
            temperature=0.7,  # 提高温度获得更多样化输出
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id,
            early_stopping=True
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
    
    # 提取assistant的回复
    if "<|im_start|>assistant" in response:
        assistant_part = response.split("<|im_start|>assistant")[-1]
        if "<|im_end|>" in assistant_part:
            return assistant_part.split("<|im_end|>")[0].strip()
    
    # 备用提取方法
    if "assistant" in response:
        return response.split("assistant")[-1].strip()
    
    return response

# 8. 综合测试
print("\n" + "="*60)
print("优化后模型测试结果")
print("="*60)

test_cases = [
    ("翻译", "Good morning → 中文"),
    ("情感分析", "我中奖了！"),
    ("情感分析", "工作丢了..."),
    ("总结", "这款手机配备最新处理器，电池续航时间长，拍照效果很好。"),
    ("写邮件", "主题：项目完成"),
    ("编程", "Python Hello World"),
    ("回答问题", "中国的首都是？"),
]

for i, (instruction, input_text) in enumerate(test_cases, 1):
    print(f"\n测试 {i}:")
    print(f"输入: {instruction}: {input_text}")
    
    try:
        result = test_optimized_model(instruction, input_text)
        print(f"输出: {result}")
    except Exception as e:
        print(f"错误: {e}")
    
    print("-" * 50)

# 9. 保存LoRA适配器（可选）
lora_path = "./qwen1.5-0.5b-lora-adapter"
model.save_pretrained(lora_path)
tokenizer.save_pretrained(lora_path)
print(f"\nLoRA适配器已保存到: {lora_path}")