LLM训练基本代码学习-文本动作提取训练思路

一、核心场景定义：文本动作提取

• 输入：注意事项数据
• 输出：固定格式、固定字段的 JSON
• 训练方式：32B 大模型生成标注 → 蒸馏到 7B
• 模型：Qwen-7B-Chat + LoRA
• 输出要求：字段不能错、格式不能乱、不能少 key
• 领域适配性（识别电力检修专属动作）。
  

二、代码

2.1测试训练

import torch from datasets import load_dataset from transformers import ( AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling ) from peft import LoraConfig, get_peft_model, TaskType import os import json # ====================== 1. 基础配置 ====================== BASE_MODEL_PATH = "Qwen/Qwen-7B-Chat" DATA_PATH = "operation_samples.jsonl" #数据集 OUTPUT_DIR = "./operation_extract_7b" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" BATCH_SIZE = 2 LEARNING_RATE = 3e-5 EPOCHS = 4 # ====================== 2. 数据预处理 ====================== tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token def preprocess_function(examples): texts = [] for input_text, output_text in zip(examples["input"], examples["output"]): # 【关键】固定指令，让模型学习结构化抽取 prompt = f"""<|im_start|>user 从以下电网操作票中提取操作步骤，输出严格为JSON格式，不要多余说明： 操作票内容：{input_text}<|im_end|> <|im_start|>assistant {output_text}<|im_end|>""" texts.append(prompt) tokenized = tokenizer( texts, truncation=True, max_length=1024, padding="max_length", return_tensors="pt"#输出tensor字典 ) tokenized["labels"] = tokenized["input_ids"].clone() return tokenized # 加载数据 dataset = load_dataset("json", data_files=DATA_PATH)["train"] tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names) data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) #把padding填充的数据设置为-100，避免训练到后面的数据 # ====================== 3. LoRA配置 ====================== model = AutoModelForCausalLM.from_pretrained( BASE_MODEL_PATH, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True ) lora_config = LoraConfig( task_type=TaskType.CAUSAL_LM, r=8,#低秩矩阵维度 lora_alpha=32,#缩放因子 target_modules=["c_attn", "c_proj"],#；loRA训练配置层 #target_modules=["c_attn", "c_proj", "w1", "w2"] lora_dropout=0.05, bias="none", ) model = get_peft_model(model, lora_config)#将loRA参数注入模型 model.print_trainable_parameters() # ====================== 4. 训练 ====================== training_args = TrainingArguments( output_dir=OUTPUT_DIR, per_device_train_batch_size=BATCH_SIZE, learning_rate=LEARNING_RATE, num_train_epochs=EPOCHS, logging_steps=10, save_steps=100, fp16=True, gradient_accumulation_steps=4, save_total_limit=2, report_to="none", ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_dataset, data_collator=data_collator, ) trainer.train() model.save_pretrained(os.path.join(OUTPUT_DIR, "final"))#参数保存，下次load tokenizer.save_pretrained(os.path.join(OUTPUT_DIR, "final")) #把当前配置好的分词器（包含词表、特殊符号、padding/truncation 规则等）完整保存到指定文件夹 # ====================== 5. 推理（专门适配你的JSON结构） ====================== def predict_operation(text): model_path = os.path.join(OUTPUT_DIR, "final")#model.loads m = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True) t = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) prompt = f"""<|im_start|>user 从以下电网操作票中提取操作步骤，输出严格为JSON格式，不要多余说明： 操作票内容：{text}<|im_end|> <|im_start|>assistant """ inputs = t(prompt, return_tensors="pt").to(DEVICE) outputs = m.generate( **inputs, max_new_tokens=512, temperature=0.05, # 极低温度，保证格式稳定 top_p=0.8, do_sample=False, eos_token_id=t.eos_token_id, ) res = t.decode(outputs[0], skip_special_tokens=True) res = res.split("assistant")[-1].strip() # 自动校验JSON，保证可用 try: return json.loads(res) except: return {"data": []} # ====================== 测试 ====================== if __name__ == "__main__": test_text = "操作A设备按正常方式复电。" result = predict_operation(test_text) print(json.dumps(result, ensure_ascii=False, indent=2))

2.2 拓展

三、基于PPO算法的再训练

「再训练过程中」，每一轮迭代实时打分，更新模型参数

import torch import json import os from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead from transformers import AutoTokenizer from peft import PeftModel # ====================== 1. 加载基础模型 ====================== BASE_MODEL_PATH = "Qwen/Qwen-7B-Chat" LORA_PATH = "./operation_extract_7b/final" # 你训练好的SFT-LoRA权重 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # 加载基础模型+LoRA权重 base_model = AutoModelForCausalLM.from_pretrained( BASE_MODEL_PATH, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True ) model = PeftModel.from_pretrained(base_model, LORA_PATH) # 包装成带价值头的模型（RL需要价值头评估奖励） model = AutoModelForCausalLMWithValueHead.from_pretrained(model, torch_dtype=torch.float16) tokenizer = AutoTokenizer.from_pretrained(LORA_PATH, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token # ====================== 2. 定义奖励函数 ====================== def calculate_reward(json_str, ground_truth): """ 计算奖励值 参数： json_str: 模型生成的JSON字符串 ground_truth: 真实标注（用于内容准确评分） 返回： total_reward: 总奖励值（0~11） """ return total_reward # ====================== 3. PPO配置 + 训练器初始化 ====================== ppo_config = PPOConfig( batch_size=1, # RL批次更小，避免显存溢出 learning_rate=1e-6, # RL学习率远低于SFT，防止模型震荡 log_with="none", mini_batch_size=1, gradient_accumulation_steps=4, ) # 初始化PPO训练器 ppo_trainer = PPOTrainer( model=model, config=ppo_config, tokenizer=tokenizer, ) # ====================== 4. RL再训练 ====================== def rl_finetune(sample): """ sample格式：{"input": , "ground_truth": 真实JSON} """ # 步骤1：构建Prompt（和SFT一致） prompt = f"""<|im_start|>user 从以下提取操作步骤，输出严格为JSON格式，不要多余说明： 内容：{sample["input"]}<|im_end|> <|im_start|>assistant """ inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE) input_ids = inputs["input_ids"] # 步骤2：模型生成JSON（当前策略） generation_kwargs = { "max_new_tokens": 512, "temperature": 0.05, "do_sample": False, } response = ppo_trainer.generate(input_ids, **generation_kwargs) response_text = tokenizer.decode(response[0], skip_special_tokens=True).split("assistant")[-1].strip() # 步骤3：计算奖励 reward = calculate_reward(response_text, sample["ground_truth"]) reward_tensor = torch.tensor([reward], device=DEVICE) # 转张量供PPO使用 # 步骤4：PPO更新模型 stats = ppo_trainer.step([input_ids[0]], [response[0]], [reward_tensor]) return stats, reward # 测试RL迭代 test_sample = { "input": "操作A站按正常方式复电。", "ground_truth": {"data": [{"seq": "1", "action": "复电", "object": {"name": "A站"}}]} } # 迭代 for i in range(10): stats, reward = rl_finetune(test_sample) print(f"第{i+1}轮RL训练，奖励值：{reward:.2f}") # ====================== 5. 保存RL优化后的模型 ====================== model.save_pretrained("./operation_extract_7b_rl") tokenizer.save_pretrained("./operation_extract_7b_rl")