import os import json import torch import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments from datasets import Dataset # 设置 Hugging Face 缓存目录 os.environ['HF_HOME'] = '/tmp/huggingface_cache' # 读取 Hugging Face Token hf_token = os.getenv('HF_TOKEN') # 请确保环境变量 HF_TOKEN 已设置 if hf_token: from huggingface_hub import HfFolder HfFolder.save_token(hf_token) else: raise ValueError("Hugging Face token 未设置") # ✅ 统一设备 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ✅ **换成兼容的学生模型** teacher_model_name = "Qwen/Qwen1.5-7B-Chat" student_model_name = "Qwen/Qwen1.5-1.8B-Chat" # ✅ **换成 Qwen1.5 1.8B 版本** # 加载教师模型 teacher = AutoModelForCausalLM.from_pretrained( teacher_model_name, trust_remote_code=True, token=hf_token ).to(device) teacher.eval() # 加载**学生模型** student = AutoModelForCausalLM.from_pretrained( student_model_name, trust_remote_code=True, token=hf_token ).to(device) tokenizer = AutoTokenizer.from_pretrained( student_model_name, # ✅ **用 Qwen 词表，防止维度不匹配** trust_remote_code=True, token=hf_token ) # **处理 pad_token 问题** if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 读取数据文件 with open('data.json', 'r', encoding='utf-8') as f: data = json.load(f) # 检查数据格式 if not isinstance(data, list): raise ValueError("data.json 格式错误，需要是一个列表！") # **格式化数据（加 `chat_template`）** def format_chat(example): instruction = example["instruction"] output = example["output"] return f"<|im_start|>system\n你是一个粉丝通软件的智能客服助手。\n<|im_end|>\n<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n{output}<|im_end|>" # **数据预处理** def preprocess_data(example): inputs = tokenizer(example["instruction"], truncation=True, padding="max_length", max_length=128) labels = tokenizer(example["output"], truncation=True, padding="max_length", max_length=128) return { "input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"], "labels": labels["input_ids"] } # def preprocess_data(example): # formatted_text = format_chat(example) # tokens = tokenizer(formatted_text, truncation=True, padding="max_length", max_length=128) # return { # "input_ids": tokens["input_ids"], # "attention_mask": tokens["attention_mask"], # "labels": tokens["input_ids"] # } dataset = Dataset.from_list(data) dataset = dataset.map(preprocess_data, batched=True) # ✅ **修正 KL 散度计算** class DistillationTrainer(Trainer): def __init__(self, teacher, *args, **kwargs): super().__init__(*args, **kwargs) self.teacher = teacher.to(device) def compute_loss(self, model, inputs, return_outputs=False): labels = inputs["input_ids"] # **学生模型的输出** outputs_student = model(**inputs) logits_student = outputs_student.logits # **教师模型的输出** with torch.no_grad(): inputs_on_device = {k: v.to(device) for k, v in inputs.items()} outputs_teacher = self.teacher(**inputs_on_device) logits_teacher = outputs_teacher.logits temperature = 2.0 # **修正维度** student_log_probs = torch.nn.functional.log_softmax(logits_student / temperature, dim=-1) teacher_probs = torch.nn.functional.softmax(logits_teacher / temperature, dim=-1) # ✅ **确保两个 logits 维度匹配** if student_log_probs.shape != teacher_probs.shape: min_dim = min(student_log_probs.shape[-1], teacher_probs.shape[-1]) student_log_probs = student_log_probs[..., :min_dim] teacher_probs = teacher_probs[..., :min_dim] # **计算 KL 散度** kl_loss = torch.nn.functional.kl_div(student_log_probs, teacher_probs, reduction="batchmean") * (temperature ** 2) # **计算交叉熵损失** ce_loss_fct = torch.nn.CrossEntropyLoss() ce_loss = ce_loss_fct(logits_student.view(-1, logits_student.size(-1)), labels.view(-1)) # **混合损失** alpha = 0.5 loss = alpha * ce_loss + (1 - alpha) * kl_loss return (loss, outputs_student) if return_outputs else loss def training_step(self, model, inputs, *args, **kwargs): """确保所有输入都在 GPU 上""" model.train() inputs = {k: v.to(device) for k, v in self._prepare_inputs(inputs).items()} loss = self.compute_loss(model, inputs) return loss # 训练参数 # ✅ 移除 `use_cache` 选项 training_args = TrainingArguments( output_dir="/tmp/distilled_model", num_train_epochs=3, per_device_train_batch_size=8, learning_rate=2e-5, weight_decay=0.01, evaluation_strategy="epoch", logging_steps=100, save_strategy="epoch", remove_unused_columns=False, gradient_checkpointing=True, fp16=False, # ✅ 禁用 fp16，避免 _scale=None 错误 bf16=True if torch.cuda.is_available() else False # ✅ 仅在支持 bf16 时启用 ) # ✅ **手动禁用 `use_cache`** student.config.use_cache = False # 🔥 这样就不会影响 `TrainingArguments`，但依然禁用了 `use_cache` # 初始化 Trainer trainer = DistillationTrainer( teacher=teacher, model=student, args=training_args, train_dataset=dataset, eval_dataset=dataset ) # 开始训练 trainer.train() # **保存模型** student.push_to_hub("Snow2222/fst-nnn", use_auth_token=hf_token) tokenizer.push_to_hub("Snow2222/fst-nnn", use_auth_token=hf_token) # ✅ 部署 Gradio Web 界面 print("🎉 训练完成，启动 Gradio Web 界面...") # **模型 ID** model_id = "Snow2222/fst-nnn" # **选择设备** device = "cuda" if torch.cuda.is_available() else "cpu" # **加载模型** print("🚀 正在加载模型...") model = AutoModelForCausalLM.from_pretrained( model_id, token=hf_token, trust_remote_code=True ).to(device) tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token) print("✅ 模型加载成功！") # **Gradio 交互函数** def chat_response(prompt): chat_input = f"<|im_start|>system\n你是一个粉丝通软件的智能客服助手。\n<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n" inputs = tokenizer(chat_input, return_tensors="pt").to(device) outputs = model.generate( **inputs, max_length=100, # ✅ 控制生成长度 do_sample=True, temperature=0.7 ) result = tokenizer.decode(outputs[0], skip_special_tokens=True) return result # **创建 Gradio 界面** iface = gr.Interface( fn=chat_response, inputs=gr.Textbox(lines=2, placeholder="请输入你的问题..."), outputs="text", title="粉丝通 AI 客服", description="基于 Snow2222/fst-nnn 训练的 AI 模型，自动回答你的问题。", allow_flagging="never" ) # **运行 Gradio** iface.launch(server_name="0.0.0.0", server_port=7860)