diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..7cb99d8
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,2 @@
+* text=auto
+* -text
diff --git a/.gitignore b/.gitignore
index cbb3920..6ccd8e4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,6 @@
 /model/__pycache__
 /dataset
 /wandb
-/out
\ No newline at end of file
+/out
+minimind-v1-small/
+minimind_dpo/
\ No newline at end of file
diff --git a/5-dpo_train.py b/5-dpo_train.py
index ef17e2c..59aed07 100644
--- a/5-dpo_train.py
+++ b/5-dpo_train.py
@@ -1,5 +1,6 @@
 import os
 import warnings
+import wandb
 
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer
@@ -12,8 +13,8 @@ warnings.filterwarnings('ignore')
 def init_model():
     device = 'cuda:0'
     # Do model patching and add fast LoRA weights
-    model_name_or_path = "minimind-v1"
-    tokenizer_name_or_path = "minimind-v1"
+    model_name_or_path = "minimind-v1-small"
+    tokenizer_name_or_path = "minimind-v1-small"
     model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True, use_fast=False)
     tokenizer.pad_token = tokenizer.eos_token
@@ -23,26 +24,56 @@ def init_model():
 
 if __name__ == '__main__':
     model, tokenizer = init_model()
+    
+    # 初始化 wandb
+    wandb.init(
+        project="minimind-dpo",
+        name=f"dpo-training-lr-4e-5-beta-0.1",
+        config={
+            "learning_rate": 4e-5,
+            "beta": 0.1,
+            "batch_size": 1,
+            "max_length": 512,
+        }
+    )
+    
     training_config = DPOConfig(
         output_dir="./minimind_dpo",
         per_device_train_batch_size=1,
         remove_unused_columns=False,
-        report_to="none",
+        report_to="wandb",  # 启用wandb记录
         save_steps=2000,
-        learning_rate=4e-5
+        learning_rate=4e-5,
+        beta=0.1,
+        max_length=512,
+        max_prompt_length=512,
+        logging_steps=50,  # 增加日志记录频率
+        logging_first_step=True  # 记录第一步
+        # 移除了 evaluation_strategy 和 eval_steps
     )
 
     dataset_path = './dataset/dpo/train_data.json'
     train_dataset = load_dataset('json', data_files=dataset_path)
+    # 将数据集转换为训练集
+    # train_dataset = train_dataset['train']
+    print("-----------------")
+    print(train_dataset)
+    import sys
+    # sys.exit(0)
 
+    # copy_model = AutoModelForCausalLM.from_pretrained("minimind-v1-small", trust_remote_code=True)
     dpo_trainer = DPOTrainer(
         model,
         ref_model=None,
         args=training_config,
-        beta=0.1,
+        # beta=0.1,
         train_dataset=train_dataset['train'],
         tokenizer=tokenizer,
-        max_length=512,
-        max_prompt_length=512
+        # max_length=512,
+        # max_prompt_length=512
     )
-    dpo_trainer.train()
+    
+    try:
+        dpo_trainer.train()
+    finally:
+        wandb.finish()  # 确保正确关闭wandb