train1.py

import os
import json
from dataclasses import dataclass
from typing import List, Dict, Any, Optional

import torch
from torch.utils.data import Dataset
import time
from datetime import datetime, timedelta

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    TrainerCallback,
    set_seed,
)
from sklearn.metrics import accuracy_score, f1_score
import numpy as np

# 尝试导入 pynvml（用于获取真实 GPU 显存占用，与 nvidia-smi 一致）
try:
    import pynvml
    PYNVML_AVAILABLE = True
except ImportError:
    PYNVML_AVAILABLE = False


# =============================================================================
# 第一部分：配置参数（你可以修改这些参数来学习它们的作用）
# =============================================================================

@dataclass
class Config:
    """
    训练配置类 - 所有可调整的参数都在这里

    【初学者指南】修改这些参数，观察训练效果的变化：

    1. learning_rate（学习率）：
       - 太大（如 0.01）：模型学习不稳定，loss 可能震荡
       - 太小（如 0.000001）：模型学习太慢
       - 推荐范围：1e-5 到 5e-5

    2. num_train_epochs（训练轮数）：
       - 太少：模型学习不充分（欠拟合）
       - 太多：模型过度学习训练数据（过拟合）
       - 推荐：先用 1-3 轮观察效果

    3. batch_size（批次大小）：
       - 受限于 GPU 显存，显存不够时需要减小
       - 可以通过 gradient_accumulation_steps 来模拟更大的 batch

    4. max_seq_length（最大序列长度）：
       - 根据数据长度设置，太长会占用更多显存
    """

    # ========== 模型配置 ==========
    model_name_or_path: str = "./model"  # 预训练模型路径

    # ========== 数据配置 ==========
    data_dir: str = "./dataset"          # 数据集目录
    max_seq_length: int = 512            # 最大序列长度（token 数量）

    # ========== 训练配置（核心参数，建议重点学习）==========
    output_dir: str = "./output_simple"  # 模型输出目录

    # 【重要】学习率 - 控制每次参数更新的步长
    # 尝试修改：2e-5, 5e-5, 1e-4，观察 loss 变化
    learning_rate: float = 1

    # 【重要】训练轮数 - 完整遍历训练集的次数
    # 尝试修改：1, 2, 3, 5，观察训练效果
    num_train_epochs: int = 1

    # 【重要】批次大小 - 每次训练使用的样本数
    # 显存不足时减小此值，如改为 1
    batch_size: int = 32

    # 梯度累积步数 - 用于模拟更大的 batch_size
    # 等效 batch_size = batch_size * gradient_accumulation_steps
    # 例如：batch_size=1, gradient_accumulation_steps=16，等效于 batch_size=16
    gradient_accumulation_steps: int = 2

    # ========== 高级配置 ==========
    weight_decay: float = 0.01           # 权重衰减（L2 正则化）
    warmup_ratio: float = 0.1            # 学习率预热比例
    max_grad_norm: float = 1.0           # 梯度裁剪阈值

    # 日志和保存
    logging_steps: int = 10              # 每多少步打印一次日志
    save_steps: int = 100                # 每多少步保存一次模型

    # 随机种子 - 保证实验可重复
    seed: int = 42


# =============================================================================
# 第二部分：Prompt 模板（定义模型的输入输出格式）
# =============================================================================

# 系统提示词 - 告诉模型它的任务是什么
SYSTEM_PROMPT = """你是游戏生成意图识别器。
任务：判断用户是否请求生成互动小游戏。
输出：1（生成游戏）或 0（拒绝生成）
只输出数字，无需解释。"""


def build_prompt(text: str, label: Optional[int] = None) -> str:
    """
    构建模型的输入 Prompt

    【学习要点】Prompt 工程
    对于大语言模型，输入格式非常重要。这里使用 ChatML 格式：
    - <|im_start|>system：系统提示，定义模型角色
    - <|im_start|>user：用户输入
    - <|im_start|>assistant：模型回复

    Args:
        text: 用户输入的文本
        label: 标签（0 或 1），训练时需要，推理时不需要

    Returns:
        格式化后的 prompt 字符串
    """
    prompt = f"""<|im_start|>system
{SYSTEM_PROMPT}<|im_end|>
<|im_start|>user
{text}<|im_end|>
<|im_start|>assistant
"""
    # 训练时添加标签，推理时不添加
    if label is not None:
        prompt += f"{label}<|im_end|>"
    return prompt


# =============================================================================
# 第三部分：数据集类（加载和处理训练数据）
# =============================================================================

class IntentDataset(Dataset):
    """
    意图识别数据集类

    【学习要点】PyTorch Dataset
    - 继承 torch.utils.data.Dataset
    - 必须实现 __len__ 和 __getitem__ 方法
    - __getitem__ 返回单个样本的张量
    """

    def __init__(self, data_path: str, tokenizer, max_seq_length: int):
        """
        初始化数据集

        Args:
            data_path: 数据文件路径（JSONL 格式）
            tokenizer: 分词器，将文本转换为 token ID
            max_seq_length: 最大序列长度
        """
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length

        # 加载数据
        self.data = []
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    item = json.loads(line)
                    self.data.append(item)

        print(f"  从 {data_path} 加载了 {len(self.data)} 条数据")

    def __len__(self) -> int:
        """返回数据集大小"""
        return len(self.data)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """
        获取第 idx 个样本

        【学习要点】样本处理流程：
        1. 获取原始文本和标签
        2. 构建 Prompt
        3. 使用 Tokenizer 转换为 token ID
        4. 创建 labels（只在答案部分计算 loss）

        Returns:
            包含 input_ids, attention_mask, labels 的字典
        """
        item = self.data[idx]
        text = item["text"]
        label = item["label"]

        # 构建完整的 prompt（包含标签）
        full_prompt = build_prompt(text, label)

        # 分词：将文本转换为 token ID 序列
        encoding = self.tokenizer(
            full_prompt,
            max_length=self.max_seq_length,
            padding=False,
            truncation=True,
            return_tensors=None,
        )

        input_ids = encoding["input_ids"]
        attention_mask = encoding["attention_mask"]

        # 创建 labels：只在答案部分计算 loss
        labels = self._create_labels(text, input_ids)

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }

    def _create_labels(self, text: str, input_ids: List[int]) -> List[int]:
        """
        创建训练标签

        【学习要点】为什么要用 -100？
        - 在计算 loss 时，-100 会被忽略
        - 我们只想让模型学习预测答案（0 或 1）
        - 不需要让模型学习生成 prompt 部分
        """
        # 获取不含答案的 prompt 长度
        prompt_without_label = build_prompt(text, label=None)
        prompt_encoding = self.tokenizer(
            prompt_without_label,
            max_length=self.max_seq_length,
            padding=False,
            truncation=True,
            return_tensors=None,
        )
        prompt_length = len(prompt_encoding["input_ids"])

        # prompt 部分设为 -100（不计算 loss），答案部分保留原始 token
        labels = [-100] * len(input_ids)
        for i in range(prompt_length, len(input_ids)):
            labels[i] = input_ids[i]

        return labels


# =============================================================================
# 第四部分：数据整理器（处理批量数据的 padding）
# =============================================================================

@dataclass
class DataCollator:
    """
    数据整理器 - 将多个样本组合成一个 batch

    【学习要点】为什么需要 Padding？
    - 不同样本的序列长度可能不同
    - 但 GPU 需要处理固定大小的张量
    - 因此需要将短序列填充（pad）到相同长度
    """

    tokenizer: Any
    max_length: int

    def __call__(self, features: List[Dict]) -> Dict[str, torch.Tensor]:
        """
        整理一个 batch 的数据

        Args:
            features: 样本列表

        Returns:
            batch 字典，包含 input_ids, attention_mask, labels
        """
        # 找到这个 batch 中的最大长度
        max_len = max(len(f["input_ids"]) for f in features)
        # 对齐到 8 的倍数（优化 GPU 计算效率）
        max_len = ((max_len + 7) // 8) * 8
        max_len = min(max_len, self.max_length)

        batch = {"input_ids": [], "attention_mask": [], "labels": []}

        for feature in features:
            input_ids = feature["input_ids"].tolist()
            attention_mask = feature["attention_mask"].tolist()
            labels = feature["labels"].tolist()

            # 计算需要填充的长度
            pad_len = max_len - len(input_ids)

            # 填充（使用 pad_token_id 填充 input_ids）
            input_ids = input_ids + [self.tokenizer.pad_token_id] * pad_len
            attention_mask = attention_mask + [0] * pad_len  # 填充位置的 attention 为 0
            labels = labels + [-100] * pad_len               # 填充位置不计算 loss

            batch["input_ids"].append(input_ids)
            batch["attention_mask"].append(attention_mask)
            batch["labels"].append(labels)

        # 转换为 PyTorch 张量
        batch["input_ids"] = torch.tensor(batch["input_ids"], dtype=torch.long)
        batch["attention_mask"] = torch.tensor(batch["attention_mask"], dtype=torch.long)
        batch["labels"] = torch.tensor(batch["labels"], dtype=torch.long)

        return batch


# =============================================================================
# 第五部分：评估函数（计算模型性能指标）
# =============================================================================

def compute_metrics(eval_pred) -> Dict[str, float]:
    """
    计算评估指标

    【学习要点】常用的分类指标：
    - Accuracy（准确率）：预测正确的比例
    - F1 Score：精确率和召回率的调和平均

    Args:
        eval_pred: (logits, labels) 元组

    Returns:
        指标字典
    """
    logits, labels = eval_pred

    predictions = []
    references = []

    for i in range(labels.shape[0]):
        # 找到 label 位置（非 -100 的位置）
        label_mask = labels[i] != -100
        label_positions = np.where(label_mask)[0]

        if len(label_positions) > 0:
            pos = label_positions[0]
            # 获取预测的 token（取概率最大的）
            pred_token = np.argmax(logits[i, pos - 1, :])
            true_token = labels[i, pos]

            predictions.append(pred_token)
            references.append(true_token)

    if len(predictions) == 0:
        return {"accuracy": 0.0, "f1": 0.0}

    predictions = np.array(predictions)
    references = np.array(references)

    accuracy = accuracy_score(references, predictions)
    f1 = f1_score(references, predictions, average="macro", zero_division=0)

    return {"accuracy": accuracy, "f1": f1}


# =============================================================================
# 第六部分：辅助函数和训练回调（用于详细日志）
# =============================================================================

def format_time(seconds: float) -> str:
    """
    将秒数格式化为 HH:MM:SS 格式

    Args:
        seconds: 秒数

    Returns:
        格式化的时间字符串
    """
    if seconds < 0:
        return "00:00:00"
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    return f"{hours:02d}:{minutes:02d}:{secs:02d}"


def get_gpu_memory_info() -> tuple:
    """
    获取 GPU 显存使用信息

    Returns:
        (已用显存GB, 总显存GB, 使用百分比) 或 (None, None, None) 如果无 GPU

    【说明】
    - 优先使用 pynvml 获取真实显存占用（与 nvidia-smi 显示一致）
    - 如果 pynvml 不可用，降级使用 PyTorch 的 memory_reserved()
    """
    if torch.cuda.is_available():
        total = torch.cuda.get_device_properties(0).total_memory / 1024**3

        # 优先使用 pynvml 获取真实显存占用（与 nvidia-smi 一致）
        if PYNVML_AVAILABLE:
            try:
                pynvml.nvmlInit()
                handle = pynvml.nvmlDeviceGetHandleByIndex(0)
                info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                used = info.used / 1024**3
                pynvml.nvmlShutdown()
                percent = (used / total) * 100
                return used, total, percent
            except Exception:
                pass

        # 降级方案：使用 PyTorch API
        reserved = torch.cuda.memory_reserved() / 1024**3
        percent = (reserved / total) * 100
        return reserved, total, percent
    return None, None, None


def create_progress_bar(progress: float, width: int = 30) -> str:
    """
    创建文本进度条

    Args:
        progress: 进度百分比 (0-100)
        width: 进度条宽度（字符数）

    Returns:
        进度条字符串，如 [████████░░░░░░░░░░░░]
    """
    filled = int(width * progress / 100)
    empty = width - filled
    return f"[{'█' * filled}{'░' * empty}]"


class DetailedLoggingCallback(TrainerCallback):
    """
    详细训练日志回调类

    【学习要点】TrainerCallback
    - Hugging Face Transformers 提供的回调机制
    - 可以在训练的不同阶段插入自定义逻辑
    - 常用回调点：on_train_begin, on_log, on_epoch_end, on_train_end
    """

    def __init__(self):
        """初始化回调"""
        self.train_start_time = None
        self.last_log_time = None
        self.total_steps = 0
        self.samples_per_step = 0
        # 关键节点记录
        self.checkpoints = [0, 25, 50, 75, 100]  # 要记录的进度百分比（包含起始点）
        self.checkpoint_printed = set()          # 已打印的节点
        self.checkpoint_data = {}                # 存储节点数据
        self.first_log = True                    # 标记是否是第一次记录日志

    def on_train_begin(self, args, state, control, **kwargs):
        """
        训练开始时调用

        【作用】记录开始时间，计算总步数
        """
        self.train_start_time = time.time()
        self.last_log_time = self.train_start_time
        self.total_steps = state.max_steps
        # 计算每步处理的样本数 = batch_size * gradient_accumulation_steps
        self.samples_per_step = args.per_device_train_batch_size * args.gradient_accumulation_steps

        # 计算关键节点对应的步数
        self.checkpoint_steps = {
            pct: int(self.total_steps * pct / 100) for pct in self.checkpoints
        }

        # 打印训练概览
        print("\n" + "=" * 70)
        print("                        训练概览")
        print("=" * 70)
        print(f"  总训练步数: {self.total_steps}")
        print(f"  每步样本数: {self.samples_per_step}")
        print(f"  开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"  关键节点: 0%(初始), 25%={self.checkpoint_steps[25]}步, 50%={self.checkpoint_steps[50]}步, "
              f"75%={self.checkpoint_steps[75]}步, 100%={self.checkpoint_steps[100]}步")
        print("=" * 70 + "\n")

    def on_log(self, args, state, control, logs=None, **kwargs):
        """
        每次记录日志时调用

        【作用】显示详细的训练进度信息
        """
        if logs is None or state.global_step == 0:
            return

        current_time = time.time()
        elapsed_time = current_time - self.train_start_time

        # 计算进度
        progress = (state.global_step / self.total_steps) * 100 if self.total_steps > 0 else 0

        # 计算训练速度（每秒样本数）
        samples_processed = state.global_step * self.samples_per_step
        samples_per_sec = samples_processed / elapsed_time if elapsed_time > 0 else 0

        # 计算预计剩余时间
        if state.global_step > 0:
            time_per_step = elapsed_time / state.global_step
            remaining_steps = self.total_steps - state.global_step
            remaining_time = time_per_step * remaining_steps
        else:
            remaining_time = 0

        # 获取 GPU 显存信息
        gpu_used, gpu_total, gpu_percent = get_gpu_memory_info()

        # 获取当前 loss 和学习率
        current_loss = logs.get("loss", 0)
        current_lr = logs.get("learning_rate", 0)
        current_epoch = logs.get("epoch", 0)

        # 第一次记录日志时，保存为 0% 节点（初始状态）
        if self.first_log and 0 not in self.checkpoint_printed:
            self.first_log = False
            self.checkpoint_printed.add(0)
            self.checkpoint_data[0] = {
                "step": state.global_step,
                "loss": current_loss,
                "lr": current_lr,
                "time": elapsed_time
            }
            print("\n" + "*" * 70)
            print(f"  ★ 关键节点 [0%] 初始状态记录")
            print("*" * 70)
            print(f"  步数: {state.global_step}/{self.total_steps}")
            print(f"  Loss: {current_loss:.4f} (初始 Loss)")
            print(f"  学习率: {current_lr:.2e}")
            print("*" * 70 + "\n")

        # 打印详细信息
        print("\n" + "=" * 70)
        print(f"                    训练进度 [Step {state.global_step}/{self.total_steps}]")
        print("=" * 70)

        # 进度条
        progress_bar = create_progress_bar(progress)
        print(f"  进度: {progress_bar} {progress:.1f}%")

        # 当前状态
        print(f"\n  当前状态:")
        print(f"    Loss: {current_loss:.4f}          学习率: {current_lr:.2e}")
        print(f"    Epoch: {current_epoch:.2f}           步数: {state.global_step}/{self.total_steps}")

        # 性能指标
        print(f"\n  性能指标:")
        print(f"    训练速度: {samples_per_sec:.2f} samples/sec")
        if gpu_used is not None:
            print(f"    GPU 显存: {gpu_used:.2f} GB / {gpu_total:.2f} GB ({gpu_percent:.1f}%)")
        else:
            print(f"    设备: CPU 或 MPS")

        # 时间估算
        print(f"\n  时间估算:")
        print(f"    已用时间: {format_time(elapsed_time)}")
        print(f"    预计剩余: {format_time(remaining_time)}")

        print("=" * 70)

        # 检查是否到达关键节点
        self._check_and_print_checkpoint(state.global_step, current_loss, current_lr, elapsed_time)

        self.last_log_time = current_time

    def _check_and_print_checkpoint(self, current_step, loss, lr, elapsed_time):
        """
        检查并打印关键节点信息

        【作用】在 25%、50%、75%、100% 进度时打印汇总信息，方便记录到实验表格
        """
        for pct in self.checkpoints:
            target_step = self.checkpoint_steps[pct]
            # 当达到或超过目标步数，且还未打印过该节点
            if current_step >= target_step and pct not in self.checkpoint_printed:
                self.checkpoint_printed.add(pct)
                # 存储节点数据
                self.checkpoint_data[pct] = {
                    "step": current_step,
                    "loss": loss,
                    "lr": lr,
                    "time": elapsed_time
                }
                # 打印关键节点信息
                print("\n" + "*" * 70)
                print(f"  ★ 关键节点 [{pct}%] 已到达!")
                print("*" * 70)
                print(f"  步数: {current_step}/{self.total_steps}")
                print(f"  Loss: {loss:.4f}")
                print(f"  学习率: {lr:.2e}")
                print(f"  已用时间: {format_time(elapsed_time)}")
                print("*" * 70 + "\n")

    def on_epoch_end(self, args, state, control, **kwargs):
        """
        每个 epoch 结束时调用

        【作用】打印 epoch 统计信息
        """
        current_time = time.time()
        elapsed_time = current_time - self.train_start_time

        print("\n" + "+" * 70)
        print(f"  Epoch {int(state.epoch)} 完成!")
        print(f"  累计用时: {format_time(elapsed_time)}")
        print("+" * 70 + "\n")

    def on_train_end(self, args, state, control, **kwargs):
        """
        训练结束时调用

        【作用】打印训练总结
        """
        total_time = time.time() - self.train_start_time
        total_samples = state.global_step * self.samples_per_step
        avg_speed = total_samples / total_time if total_time > 0 else 0

        print("\n" + "=" * 70)
        print("                        训练完成!")
        print("=" * 70)
        print(f"  总训练时间: {format_time(total_time)}")
        print(f"  总训练步数: {state.global_step}")
        print(f"  处理样本数: {total_samples}")
        print(f"  平均速度: {avg_speed:.2f} samples/sec")
        print(f"  结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("=" * 70)

        # 打印关键节点汇总表（方便复制到实验记录）
        print("\n" + "=" * 70)
        print("                  关键节点汇总 (可复制到实验记录)")
        print("=" * 70)
        print(f"  {'节点':<8} {'步数':<12} {'Loss':<12} {'学习率':<14} {'用时':<10}")
        print("-" * 70)

        for pct in self.checkpoints:
            if pct in self.checkpoint_data:
                data = self.checkpoint_data[pct]
                print(f"  {pct}%{'':<5} {data['step']:<12} {data['loss']:<12.4f} {data['lr']:<14.2e} {format_time(data['time']):<10}")
            else:
                print(f"  {pct}%{'':<5} {'--':<12} {'--':<12} {'--':<14} {'--':<10}")

        print("=" * 70)

        # 打印简洁格式（方便直接复制数值）
        print("\n  📋 快速复制格式:")
        losses = [f"{self.checkpoint_data.get(pct, {}).get('loss', '--'):.4f}" if pct in self.checkpoint_data else "--"
                  for pct in [0, 25, 50, 75, 100]]
        print(f"  Loss@0%: {losses[0]} | Loss@25%: {losses[1]} | Loss@50%: {losses[2]} | Loss@75%: {losses[3]} | Loss@100%: {losses[4]}")
        print(f"  总时间: {format_time(total_time)}")
        print("=" * 70 + "\n")


# =============================================================================
# 第七部分：主训练函数
# =============================================================================

def main():
    """
    主训练函数 - 训练流程的核心

    【训练流程概述】
    1. 加载配置
    2. 加载分词器（Tokenizer）
    3. 加载预训练模型
    4. 加载数据集
    5. 配置训练参数
    6. 创建 Trainer 并开始训练
    7. 保存模型
    """

    print("=" * 70)
    print("         初学者友好的模型训练教程 - 开始训练")
    print("=" * 70)

    # ----- 步骤 1：加载配置 -----
    config = Config()

    # 设置随机种子，保证实验可重复
    set_seed(config.seed)

    print("\n📋 训练配置:")
    print(f"  模型路径: {config.model_name_or_path}")
    print(f"  数据目录: {config.data_dir}")
    print(f"  输出目录: {config.output_dir}")
    print(f"  学习率: {config.learning_rate}")
    print(f"  训练轮数: {config.num_train_epochs}")
    print(f"  批次大小: {config.batch_size}")
    print(f"  梯度累积: {config.gradient_accumulation_steps}")
    print(f"  等效批次大小: {config.batch_size * config.gradient_accumulation_steps}")

    # ----- 步骤 2：检查设备 -----
    print("\n🖥️ 设备信息:")
    if torch.cuda.is_available():
        print(f"  使用 GPU: {torch.cuda.get_device_name(0)}")
        print(f"  显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
        torch.cuda.empty_cache()
    elif torch.backends.mps.is_available():
        print("  使用 Apple MPS (Metal Performance Shaders)")
    else:
        print("  使用 CPU（训练会较慢）")

    # ----- 步骤 3：加载分词器 -----
    print("\n📚 加载分词器...")
    tokenizer = AutoTokenizer.from_pretrained(
        config.model_name_or_path,
        trust_remote_code=True,
        padding_side="right",
    )

    # 确保有 pad token
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    print(f"  词汇表大小: {tokenizer.vocab_size}")

    # ----- 步骤 4：加载预训练模型 -----
    print("\n🤖 加载预训练模型...")

    # 使用全精度训练 (float32)，避免混合精度兼容性问题
    if torch.cuda.is_available():
        device_map = "auto"
    else:
        device_map = None

    model = AutoModelForCausalLM.from_pretrained(
        config.model_name_or_path,
        trust_remote_code=True,
        torch_dtype=torch.float32,  # 全精度训练
        device_map=device_map,
    )

    # 关闭缓存以节省显存
    if hasattr(model, "config"):
        model.config.use_cache = False

    print(f"  模型参数量: {model.num_parameters():,}")

    # ----- 步骤 5：加载数据集 -----
    print("\n📊 加载数据集...")

    train_dataset = IntentDataset(
        data_path=os.path.join(config.data_dir, "train.jsonl"),
        tokenizer=tokenizer,
        max_seq_length=config.max_seq_length,
    )

    valid_dataset = IntentDataset(
        data_path=os.path.join(config.data_dir, "valid.jsonl"),
        tokenizer=tokenizer,
        max_seq_length=config.max_seq_length,
    )

    # 创建数据整理器
    data_collator = DataCollator(
        tokenizer=tokenizer,
        max_length=config.max_seq_length,
    )

    # ----- 步骤 6：配置训练参数 -----
    print("\n⚙️ 配置训练参数...")

    # 创建输出目录
    os.makedirs(config.output_dir, exist_ok=True)

    training_args = TrainingArguments(
        # 输出目录
        output_dir=config.output_dir,

        # 核心训练参数
        num_train_epochs=config.num_train_epochs,
        per_device_train_batch_size=config.batch_size,
        per_device_eval_batch_size=config.batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,

        # 学习率相关
        learning_rate=config.learning_rate,
        weight_decay=config.weight_decay,
        warmup_ratio=config.warmup_ratio,
        lr_scheduler_type="cosine",  # 学习率调度器：cosine 衰减

        # 优化器
        optim="adamw_torch",
        max_grad_norm=config.max_grad_norm,

        # 日志和保存
        logging_steps=config.logging_steps,
        save_steps=config.save_steps,
        save_total_limit=2,  # 最多保存 2 个 checkpoint

        # 评估（暂时关闭以加快训练）
        eval_strategy="no",

        # 混合精度 - 使用全精度训练，避免兼容性问题
        fp16=False,
        bf16=False,

        # 梯度检查点 - 关闭以避免与优化器的兼容问题
        gradient_checkpointing=False,

        # 其他
        dataloader_num_workers=0,
        seed=config.seed,
        report_to="none",  # 不使用 tensorboard 等报告工具
        remove_unused_columns=False,
        logging_dir=os.path.join(config.output_dir, "logs"),
    )

    # ----- 步骤 7：创建 Trainer -----
    print("\n🏋️ 创建 Trainer...")

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[DetailedLoggingCallback()],  # 添加详细日志回调
    )

    # ----- 步骤 8：开始训练 -----
    print("\n【训练日志说明】")
    print("  Loss: 损失值，越小越好，表示模型预测与真实标签的差距")
    print("  学习率: 控制每次参数更新的步长，会随训练逐渐衰减")
    print("  训练速度: 每秒处理的样本数量")
    print("  GPU 显存: 当前显存占用情况")

    trainer.train()

    # ----- 步骤 9：保存最终模型 -----
    print("\n💾 保存模型...")
    final_model_path = os.path.join(config.output_dir, "final_model")
    trainer.save_model(final_model_path)
    tokenizer.save_pretrained(final_model_path)

    print("\n" + "=" * 70)
    print(f"训练完成！模型已保存到: {final_model_path}")
    print("=" * 70)

    # ----- 可选：在验证集上评估 -----
    print("\n📈 在验证集上评估...")
    eval_results = trainer.evaluate()
    print("评估结果:")
    for key, value in eval_results.items():
        if isinstance(value, float):
            print(f"  {key}: {value:.4f}")


# =============================================================================
# 第七部分：运行入口
# =============================================================================

if __name__ == "__main__":
    main()