frontend/train_page.py

@@ -46,10 +46,30 @@ def train_page():
             dataset = get_datasets().get(Query().name == dataset_name)
             dataset = [ds["message"][0] for ds in dataset["dataset_items"]] 
             
+            class LossCallback(TrainerCallback):
+                def __init__(self):
+                    self.loss_data = []
+                    self.log_text = ""
+                    self.last_output = {"text": "", "plot": None}
+                def on_log(self, args, state, control, logs=None, **kwargs):
+                    if "loss" in logs:
+                        self.loss_data.append({
+                            "step": state.global_step,
+                            "loss": float(logs["loss"])
+                        })
+                    self.log_text += f"Step {state.global_step}: loss={logs['loss']:.4f}\n"
+                    # 添加以下两行print语句
+                    print(f"Current Step: {state.global_step}")
+                    print(f"Loss Value: {logs['loss']:.4f}")
+                    self.last_output = {
+                        "text": self.log_text,
+                    }
+                    # 不返回 control，避免干预训练过程
+
             train_model(get_model(), get_tokenizer(),
-                         dataset, get_workdir() + "/training" + "/1", 
+                         dataset, get_workdir()+"/checkpoint", 
                          learning_rate, per_device_train_batch_size, epoch,
-                           save_steps, lora_rank)
+                           save_steps, lora_rank, LossCallback)
 
 
         train_button.click(
@@ -60,7 +80,7 @@ def train_page():
                 per_device_train_batch_size_input,
                 epoch_input,
                 save_steps_input,
-                lora_rank_input
+                lora_rank_input  # 新增lora_rank_input
             ],
             outputs=output
         )

main.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import unsloth
+from frontend.setting_page import setting_page
 from frontend import *
 from db import initialize_sqlite_db, initialize_prompt_store
 from global_var import init_global_var, get_sql_engine, get_prompt_store

requirements.txt

@@ -7,5 +7,3 @@ tinydb>=4.0.0
 unsloth>=2025.3.19
 sqlmodel>=0.0.24
 jinja2>=3.1.0
-tensorboardX>=2.6.2.2
-tensorboard>=2.19.0

tools/model.py

@@ -35,13 +35,13 @@ def train_model(
     model, 
     tokenizer, 
     dataset: list, 
-    train_dir: str, 
+    output_dir: str, 
     learning_rate: float, 
     per_device_train_batch_size: int, 
     epoch: int, 
     save_steps: int, 
     lora_rank: int,
-    trainer_callback=None
+    trainer_callback
 ) -> None:
     # 模型配置参数
     dtype = None  # 数据类型，None表示自动选择
@@ -113,15 +113,13 @@ def train_model(
             optim="adamw_8bit",  # 使用8位AdamW优化器节省显存，几乎不影响训练效果
             weight_decay=0.01,  # 权重衰减系数,用于正则化，防止过拟合
             seed=114514,  # 随机数种子
-            output_dir=train_dir + "/checkpoints",  # 保存模型检查点和训练日志
+            output_dir=output_dir,  # 保存模型检查点和训练日志
             save_strategy="steps",  # 按步保存中间权重
             save_steps=save_steps,  # 使用动态传入的保存步数
-            logging_dir=train_dir + "/logs",  # 日志文件存储路径
-            report_to="tensorboard",  # 使用TensorBoard记录日志
+            report_to="none",
         ),
     )
 
-    if trainer_callback is not None:
     trainer.add_callback(trainer_callback)
 
     trainer = train_on_responses_only(