训练一个小语言模型(Qwen2.5-0.5B),从产品名称和特征生成有吸引力的、简洁的商业产品描述。这个项目整合了所有学到的概念:QLoRA模型加载、LoRA配置、数据集准备、SFTTrainer训练、评估和适配器保存。
创建一个包含50个示例的合成数据集,采用Alpaca格式:
import json
dataset = [
{
"instruction": "为科技产品写一个有吸引力的商业描述。",
"input": "名称:可折叠4K无人机。特性:稳定摄像头、30分钟飞行时间、应用程序控制、自动跟踪模式。价格:$349.99。",
"output": "用我们的可折叠4K无人机从上方捕捉世界!轻便、强大、随时准备旅行。得益于其稳定摄像头,享受超清晰视频,无忧飞行30分钟,通过智能手机控制一切。自动跟踪模式让您永远不会失去对行动的视线。专业科技触手可及。仅售$349.99。"
},
{
"instruction": "为科技产品写一个有吸引力的商业描述。",
"input": "名称:RGB机械键盘。特性:触觉开关、可定制背光、人体工学设计、防鬼键。价格:$89.99。",
"output": "用我们的RGB机械键盘主宰每个游戏或项目。专为那些要求精度和风格的人设计。触觉开关响应您的触摸,RGB背光让您创造理想的氛围,人体工学设计在长时间会话中保护您的双手。防鬼键确保每次按键都有效。性能和优雅于一体。仅售$89.99。"
},
# ... 还有48个示例
]
# 保存为JSON
with open("product_descriptions_dataset.json", "w") as f:
json.dump(dataset, f, indent=2)
# --- 导入 ---
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
TrainingArguments
)
from peft import LoraConfig, get_peft_model, PeftModel
from trl import SFTTrainer
from datasets import Dataset
import json
import wandb
# --- 初始设置 ---
wandb.login() # 可选
# --- 加载数据集 ---
with open("product_descriptions_dataset.json", "r") as f:
raw_data = json.load(f)
dataset = Dataset.from_list(raw_data)
# --- 使用QLoRA加载模型和Tokenizer ---
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token # 填充所需
# --- 配置LoRA ---
lora_config = LoraConfig(
r=8,
lora_alpha=16,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# --- 格式化函数 ---
def formatting_prompts_func(examples):
instructions = examples["instruction"]
inputs = examples["input"]
outputs = examples["output"]
texts = []
for instruction, input_text, output in zip(instructions, inputs, outputs):
text = f"### 指令:\n{instruction}\n\n### 输入:\n{input_text}\n\n### 响应:\n{output}"
texts.append(text)
return texts
# --- 训练配置 ---
training_args = TrainingArguments(
output_dir="./product_desc_results",
num_train_epochs=5,
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
optim="paged_adamw_8bit",
save_steps=100,
logging_steps=25,
learning_rate=2e-4,
weight_decay=0.01,
fp16=True,
max_grad_norm=0.3,
warmup_ratio=0.03,
lr_scheduler_type="cosine",
report_to="wandb",
evaluation_strategy="no", # 为简化跳过评估
save_total_limit=2,
)
# --- 创建并运行Trainer ---
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset,
formatting_func=formatting_prompts_func,
max_seq_length=512,
tokenizer=tokenizer,
packing=False,
)
trainer.train()
# --- 保存适配器 ---
model.save_pretrained("./product_desc_adapter")
tokenizer.save_pretrained("./product_desc_adapter")
# --- 可选合并并保存完整模型 ---
# model = model.merge_and_unload()
# model.save_pretrained("./product_desc_merged")
# tokenizer.save_pretrained("./product_desc_merged")
训练后,在新产品上测试模型:
def generate_product_description(product_name, features, price):
instruction = "为科技产品写一个有吸引力的商业描述。"
input_text = f"名称: {product_name}. 特性: {features}. 价格: ${price}."
prompt = f"### 指令:\n{instruction}\n\n### 输入:\n{input_text}\n\n### 响应:\n"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=200,
temperature=0.7,
top_p=0.9,
do_sample=True
)
full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
response = full_text.split("### 响应:")[-1].strip()
return response
# 测试
desc = generate_product_description(
"防水蓝牙音箱",
"20W功率、15小时电池、多点连接、便携设计",
"79.99"
)
print(desc)
预期输出(示例):
"用我们的防水蓝牙音箱随时随地享受派对!纯净的20W功率在任何空间都充满沉浸式声音。享受长达15小时的不间断音乐,同时连接两个设备,并通过其坚固便携的设计将其带到海滩、游泳池或山上。无拘无束的乐趣无论您去哪里。仅售$79.99。"