Update train.py

train.py

@@ -1,58 +1,70 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+import os
+import json
+from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer
+from datasets import Dataset
+from huggingface_hub import HfApi, HfFolder
 from peft import LoraConfig, get_peft_model
-from datasets import load_dataset
 
-# 加载DeepSeek R1模型
-model_name = "DeepSeek/R1"  # 你可以根据实际选择不同的路径
+# 从环境变量加载 Hugging Face token
+hf_token = os.getenv('HF_TOKEN')  # 假设你将 token 设置为环境变量
+if hf_token:
+    HfFolder.save_token(hf_token)
+else:
+    raise ValueError("Hugging Face token 未设置")
 
-# 加载模型和分词器
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# 加载基础模型（例如：DeepSeek-R1）
+model_name = "DeepSeek-R1"  # 你可以根据需要调整基础模型
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+# 加载 tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# 配置 LoRA 微调
-lora_config = LoraConfig(
-    r=8,  # LoRA参数
-    lora_alpha=16,
-    lora_dropout=0.1,
-    bias="none",
-)
+# 读取数据文件
+with open('data.json', 'r', encoding='utf-8') as f:
+    data = json.load(f)
 
-# 获取LoRA微调模型
-model = get_peft_model(model, lora_config)
+# 将数据转换为 Dataset 格式
+def preprocess_data(example):
+    return tokenizer(example['instruction'], truncation=True, padding="max_length", max_length=128)
 
-# 准备数据集
-data = [
-    {"instruction": "粉丝通跨店版的费用是多少？", "output": "粉丝通跨店版按月付费，500元/月，仅提供增值税普通电子发票。"},
-    {"instruction": "如何充值粉丝通软件的红包？", "output": "商家可以灵活充值红包，每个红包最低0.1元，具体总额根据拉新目标决定。"},
-    {"instruction": "红包的扣费机制是怎样的？", "output": "红包在用户实际使用后才会扣款，未使用到期会自动退回商家公户。"},
-    # 你可以继续添加数据...
-]
+dataset = Dataset.from_dict(data)
+dataset = dataset.map(preprocess_data, batched=True)
 
-# 转换数据为 Hugging Face 数据集格式
-train_data = [{"input_ids": tokenizer.encode(d["instruction"], truncation=True, padding="max_length"), "labels": tokenizer.encode(d["output"], truncation=True, padding="max_length")} for d in data]
+# LoRA 配置（根据需要调整）
+lora_config = LoraConfig(
+    r=8,
+    lora_alpha=32,
+    lora_dropout=0.1,
+    target_modules=["q_proj", "v_proj"],  # 根据你的模型结构调整
+)
 
-train_dataset = load_dataset('json', data_files={'train': train_data})
+# 获取 LoRA 微调模型
+peft_model = get_peft_model(model, lora_config)
 
 # 设置训练参数
 training_args = TrainingArguments(
-    output_dir='./results',
+    output_dir="./output",
     evaluation_strategy="epoch",
-    learning_rate=5e-5,
-    per_device_train_batch_size=2,
-    per_device_eval_batch_size=2,
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
     num_train_epochs=3,
     weight_decay=0.01,
-    save_steps=10_000,
-    save_total_limit=2,
+    logging_dir="./logs",
 )
 
-# 设置 Trainer
+# 训练器
 trainer = Trainer(
-    model=model,
+    model=peft_model,
     args=training_args,
-    train_dataset=train_dataset["train"],
+    train_dataset=dataset,
+    eval_dataset=dataset,  # 你可以根据需要设置验证集
 )
 
 # 开始训练
 trainer.train()
+
+# 保存模型
+model.save_pretrained("./fst-nnn")
+tokenizer.save_pretrained("./fst-nnn")
+
+# 上传到 Hugging Face