1. 数据准备与预处理
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from datasets import Dataset
def load_and_preprocess_data(file_path):
"""
加载并预处理对话数据
参数:
file_path: 数据文件路径,支持csv、json格式
返回:
train_dataset, val_dataset: 训练集和验证集
"""
# 读取数据
if file_path.endswith('.csv'):
df = pd.read_csv(file_path)
elif file_path.endswith('.json'):
df = pd.read_json(file_path)
else:
raise ValueError("不支持的数据格式,请使用csv或json")
# 确保数据包含必要的列
required_columns = ['prompt', 'response']
if not all(col in df.columns for col in required_columns):
raise ValueError(f"数据必须包含以下列: {required_columns}")
# 重命名列以匹配我们的提示词格式
df = df.rename(columns={'prompt': 'user_message', 'response': 'assistant_message'})
# 构建完整的对话示例
df['full_prompt'] = df.apply(
lambda row: f"<|SYSTEM|># StableLM对话助手\n你是一个乐于助人的AI助手。\n<|USER|>{row['user_message']}<|ASSISTANT|>{row['assistant_message']}",
axis=1
)
# 划分训练集和验证集
train_df, val_df = train_test_split(df, test_size=0.1, random_state=42)
# 转换为Hugging Face Dataset格式
train_dataset = Dataset.from_pandas(train_df)
val_dataset = Dataset.from_pandas(val_df)
return train_dataset, val_dataset
2. 模型微调代码
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
import torch
def load_and_prepare_model(model_name_or_path):
"""
加载预训练模型和分词器
参数:
model_name_or_path: 模型路径或名称
返回:
model: 加载的模型
tokenizer: 对应的分词器
"""
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
# 如果模型没有pad_token,设置为eos_token
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# 加载模型,使用4位量化以节省显存
model = AutoModelForCausalLM.from_pretrained(
model_name_or_path,
load_in_4bit=True,
device_map="auto",
torch_dtype=torch.float16
)
return model, tokenizer
def tokenize_function(examples, tokenizer, max_length=2048):
"""
对文本进行分词处理
参数:
examples: 待分词的文本数据
tokenizer: 使用的分词器
max_length: 最大序列长度
返回:
分词后的结果
"""
# 对输入文本进行分词
return tokenizer(
examples["full_prompt"],
truncation=True,
max_length=max_length,
padding="max_length",
return_tensors="pt"
)
def fine_tune_model(
model,
tokenizer,
train_dataset,
val_dataset,
output_dir="./stablelm-finetuned",
max_steps=1000,
batch_size=4,
learning_rate=2e-5
):
"""
微调模型
参数:
model: 基础模型
tokenizer: 分词器
train_dataset: 训练数据集
val_dataset: 验证数据集
output_dir: 模型保存目录
max_steps: 训练最大步数
batch_size: 批处理大小
learning_rate: 学习率
返回:
trainer: 训练器对象
"""
# 对数据集进行分词
tokenized_train = train_dataset.map(
lambda x: tokenize_function(x, tokenizer),
batched=True,
remove_columns=train_dataset.column_names
)
tokenized_val = val_dataset.map(
lambda x: tokenize_function(x, tokenizer),
batched=True,
remove_columns=val_dataset.column_names
)
# 数据收集器,用于语言模型训练
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False, # 因果语言模型不需要掩码语言模型
)
# 训练参数设置
training_args = TrainingArguments(
output_dir=output_dir,
num_train_epochs=1, # 实际使用中可能需要调整为多个epochs
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
gradient_accumulation_steps=4,
learning_rate=learning_rate,
weight_decay=0.01,
logging_dir=f"{output_dir}/logs",
logging_steps=10,
evaluation_strategy="steps",
eval_steps=50,
save_strategy="steps",
save_steps=100,
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
fp16=True, # 使用混合精度训练
report_to="tensorboard",
max_steps=max_steps
)
# 创建Trainer对象
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_train,
eval_dataset=tokenized_val,
data_collator=data_collator,
)
# 开始训练
trainer.train()
return trainer
3. 微调后模型加载与评估
def load_finetuned_model(model_dir):
"""
加载微调后的模型
参数:
model_dir: 微调后模型保存目录
返回:
model: 微调后的模型
tokenizer: 对应的分词器
"""
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_dir)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# 加载模型
model = AutoModelForCausalLM.from_pretrained(
model_dir,
load_in_4bit=True,
device_map="auto",
torch_dtype=torch.float16
)
return model, tokenizer
def evaluate_finetuned_model(model, tokenizer, test_dataset, max_length=2048):
"""
评估微调后的模型性能
参数:
model: 微调后的模型
tokenizer: 分词器
test_dataset: 测试数据集
max_length: 最大序列长度
返回:
results: 评估结果
"""
from evaluate import load
bleu = load("bleu")
rouge = load("rouge")
results = {"bleu": [], "rouge": []}
# 随机选择一些样本进行评估
test_samples = test_dataset.shuffle(seed=42).select(range(min(100, len(test_dataset))))
for sample in test_samples:
prompt = sample["user_message"]
# 构建完整提示词
full_prompt = f"<|SYSTEM|># StableLM对话助手\n你是一个乐于助人的AI助手。\n<|USER|>{prompt}<|ASSISTANT|>"
# 分词和生成
inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=max_length).to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=128,
temperature=0.7,
do_sample=True
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|ASSISTANT|>")[-1].strip()
# 计算评估指标
reference = sample["assistant_message"]
# BLEU计算
bleu_score = bleu.compute(predictions=[generated_text], references=[[reference]])
results["bleu"].append(bleu_score["bleu"])
# ROUGE计算
rouge_score = rouge.compute(predictions=[generated_text], references=[reference])
results["rouge"].append(rouge_score)
# 计算平均指标
avg_bleu = sum(results["bleu"]) / len(results["bleu"])
avg_rouge = {key: sum([score[key] for score in results["rouge"]]) / len(results["rouge"]) for key in results["rouge"][0]}
print(f"平均BLEU分数: {avg_bleu:.4f}")
print(f"平均ROUGE分数: {avg_rouge}")
return {
"avg_bleu": avg_bleu,
"avg_rouge": avg_rouge,
"sample_count": len(test_samples)
}
4. 全流程使用示例
if __name__ == "__main__":
# 1. 数据准备
print("===== 1. 数据准备 =====")
train_dataset, val_dataset = load_and_preprocess_data("your_dialogue_data.csv") # 替换为实际数据路径
print(f"训练集大小: {len(train_dataset)}, 验证集大小: {len(val_dataset)}")
# 2. 加载基础模型和分词器
print("\n===== 2. 加载基础模型 =====")
model, tokenizer = load_and_prepare_model("01-stablelm-tuned-alpha-7b") # 替换为基础模型路径
# 3. 模型微调
print("\n===== 3. 模型微调 =====")
trainer = fine_tune_model(
model,
tokenizer,
train_dataset,
val_dataset,
output_dir="./stablelm-finetuned",
max_steps=2000,
batch_size=4,
learning_rate=2e-5
)
# 4. 评估微调后的模型
print("\n===== 4. 模型评估 =====")
# 假设我们有一个测试集
test_dataset = val_dataset.shuffle(seed=42).select(range(100)) # 使用验证集的一部分作为测试集
evaluation_results = evaluate_finetuned_model(model, tokenizer, test_dataset)
# 5. 加载微调后的模型并测试对话
print("\n===== 5. 测试对话 =====")
finetuned_model, finetuned_tokenizer = load_finetuned_model("./stablelm-finetuned")
def chat_with_finetuned_model(prompt, max_new_tokens=256, temperature=0.7):
system_prompt = """<|SYSTEM|># StableLM对话助手
你是一个乐于助人的AI助手,能够根据用户的问题提供准确、有用的回答。
"""
full_prompt = f"{system_prompt}<|USER|>{prompt}<|ASSISTANT|>"
inputs = finetuned_tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=2048).to(finetuned_model.device)
with torch.no_grad():
outputs = finetuned_model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
do_sample=True
)
response = finetuned_tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|ASSISTANT|>")[-1].strip()
return response
# 测试对话
while True:
user_input = input("\n请输入您的问题: ")
if user_input.lower() in ["exit", "quit"]:
break
response = chat_with_finetuned_model(user_input)
print(f"助手: {response}")
5. 关键注意事项
-
数据质量:微调效果高度依赖训练数据质量,建议使用干净、多样化且与目标任务相关的对话数据。
-
硬件要求:4位量化模型仍需要至少10GB以上显存,8GB显存可能勉强运行但速度较慢。
-
超参数调整:learning_rate(学习率)、batch_size(批大小)和max_steps(训练步数)需要根据实际情况调整,避免过拟合或欠拟合。
-
训练监控:建议使用TensorBoard监控训练过程中的损失值变化,及时发现过拟合或训练不收敛问题。
-
模型保存与加载:使用
trainer.save_model()保存模型,使用load_finetuned_model()加载,确保分词器和模型路径正确。 -
安全与合规:训练数据需确保无隐私泄露风险,生成内容需符合伦理规范和法律法规。
通过上述步骤,你可以基于StableLM-Tuned-Alpha-7B模型进行领域微调,使其更好地适配特定应用场景,如客服问答、技术支持、教育辅导等。
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考



