Update train.py

train.py

@@ -1,27 +1,154 @@
 import json
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    TrainingArguments, 
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from datasets import Dataset
+import os
 
 def load_training_data():
-    with open('train_data.json', 'r', encoding='utf-8') as f:
-        return json.load(f)
+    """加载训练数据"""
+    try:
+        with open('train_data.json', 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        print(f"📊 加载了 {len(data)} 条训练数据")
+        return data
+    except FileNotFoundError:
+        print("❌ 训练数据文件不存在，使用示例数据")
+        # 返回一些示例数据
+        return [
+            {
+                "input": "print('hello",
+                "output": "print('hello')",
+                "language": "python"
+            },
+            {
+                "input": "<div class=test>",
+                "output": "<div class=\"test\">",
+                "language": "html"
+            }
+        ]
+
+def prepare_dataset(data):
+    """准备训练数据集"""
+    texts = []
+    
+    for item in data:
+        # 创建训练文本格式
+        prompt = f"修复以下{item.get('language', 'code')}代码:\n{item['input']}\n修复后:\n{item['output']}"
+        texts.append(prompt)
+    
+    return Dataset.from_dict({"text": texts})
 
 def train_model():
+    """训练模型"""
+    print("🚀 开始训练代码修复模型...")
+    
     # 加载数据
     data = load_training_data()
     
-    # 初始化tokenizer和模型（使用小模型）
-    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
-    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
+    if len(data) < 5:
+        print("❌ 训练数据不足，至少需要5条数据")
+        return
+    
+    # 初始化模型和分词器
+    model_name = "microsoft/DialoGPT-small"
+    
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        
+        # 添加pad token如果不存在
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # 准备数据集
+        dataset = prepare_dataset(data)
+        
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"], 
+                truncation=True, 
+                padding=True, 
+                max_length=512
+            )
+        
+        tokenized_dataset = dataset.map(tokenize_function, batched=True)
+        
+        # 训练参数
+        training_args = TrainingArguments(
+            output_dir="./codefix-model",
+            overwrite_output_dir=True,
+            num_train_epochs=3,
+            per_device_train_batch_size=2,
+            save_steps=500,
+            save_total_limit=2,
+            logging_steps=100,
+            prediction_loss_only=True,
+            remove_unused_columns=False,
+        )
+        
+        # 数据收集器
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False,  # 不使用掩码语言模型
+        )
+        
+        # 训练器
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            data_collator=data_collator,
+            train_dataset=tokenized_dataset,
+        )
+        
+        # 开始训练
+        print("🔥 开始模型训练...")
+        trainer.train()
+        
+        # 保存模型
+        trainer.save_model()
+        tokenizer.save_pretrained("./codefix-model")
+        
+        print("✅ 模型训练完成！保存在 ./codefix-model 目录")
+        
+    except Exception as e:
+        print(f"❌ 训练失败: {e}")
+
+def incremental_train(new_feedback_file="user_feedback.json"):
+    """增量训练 - 基于用户反馈"""
+    if not os.path.exists(new_feedback_file):
+        print("❌ 用户反馈文件不存在")
+        return
     
-    # 数据处理和训练逻辑
-    # ... 这里添加你的训练代码
+    with open(new_feedback_file, 'r', encoding='utf-8') as f:
+        feedback_data = json.load(f)
     
-    # 保存模型
-    model.save_pretrained("./trained_model")
-    tokenizer.save_pretrained("./trained_model")
+    # 只使用正确的反馈作为训练数据
+    training_data = []
+    for feedback in feedback_data:
+        if feedback.get("correct", False):
+            training_data.append({
+                "input": feedback["original"],
+                "output": feedback["fixed"],
+                "language": feedback["language"]
+            })
     
-    print("模型训练完成！")
+    if len(training_data) > 0:
+        print(f"🔄 基于 {len(training_data)} 条用户反馈进行增量训练")
+        # 这里可以调用训练函数进行增量训练
+        # 为了简化，暂时只���录
+        print("📝 增量训练数据已准备就绪")
 
 if __name__ == "__main__":
-    train_model()
+    # 检查是否进行增量训练
+    import sys
+    if len(sys.argv) > 1 and sys.argv[1] == "incremental":
+        incremental_train()
+    else:
+        train_model()
+