train.py

"""
训练脚本 - Qwen 意图识别模型全参数微调
"""

import os
import argparse
from typing import Dict, Tuple

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

from config import (
    ModelConfig,
    DataConfig,
    TrainingConfig,
    get_default_configs,
)
from data_utils import IntentDataset, IntentDataCollator, load_datasets


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="训练 Qwen 意图识别模型")

    # 模型相关
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default=None,
        help="模型路径或 Hugging Face 模型名",
    )

    # 数据相关
    parser.add_argument(
        "--data_dir",
        type=str,
        default=None,
        help="数据集目录",
    )
    parser.add_argument(
        "--max_seq_length",
        type=int,
        default=None,
        help="最大序列长度",
    )

    # 训练相关
    parser.add_argument(
        "--output_dir",
        type=str,
        default=None,
        help="输出目录",
    )
    parser.add_argument(
        "--num_train_epochs",
        type=int,
        default=None,
        help="训练轮数",
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=None,
        help="每个设备的训练 batch size",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=None,
        help="学习率",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=None,
        help="梯度累积步数",
    )
    parser.add_argument(
        "--fp16",
        action="store_true",
        help="使用 FP16 混合精度训练",
    )
    parser.add_argument(
        "--bf16",
        action="store_true",
        help="使用 BF16 混合精度训练",
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help="从 checkpoint 恢复训练",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="随机种子",
    )

    return parser.parse_args()


def update_configs_from_args(
    args,
    model_config: ModelConfig,
    data_config: DataConfig,
    training_config: TrainingConfig,
) -> Tuple[ModelConfig, DataConfig, TrainingConfig]:
    """根据命令行参数更新配置"""

    # 更新模型配置
    if args.model_name_or_path is not None:
        model_config.model_name_or_path = args.model_name_or_path

    # 更新数据配置
    if args.data_dir is not None:
        data_config.data_dir = args.data_dir
    if args.max_seq_length is not None:
        data_config.max_seq_length = args.max_seq_length

    # 更新训练配置
    if args.output_dir is not None:
        training_config.output_dir = args.output_dir
    if args.num_train_epochs is not None:
        training_config.num_train_epochs = args.num_train_epochs
    if args.per_device_train_batch_size is not None:
        training_config.per_device_train_batch_size = args.per_device_train_batch_size
    if args.learning_rate is not None:
        training_config.learning_rate = args.learning_rate
    if args.gradient_accumulation_steps is not None:
        training_config.gradient_accumulation_steps = args.gradient_accumulation_steps
    if args.fp16:
        training_config.fp16 = True
        training_config.bf16 = False
    if args.bf16:
        training_config.bf16 = True
        training_config.fp16 = False
    if args.seed is not None:
        training_config.seed = args.seed

    return model_config, data_config, training_config


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

    对于生成式模型，需要从 logits 中提取预测的 label

    Args:
        eval_pred: (logits, labels) 元组

    Returns:
        包含各项指标的字典
    """
    logits, labels = eval_pred

    # logits 形状: (batch_size, seq_length, vocab_size)
    # labels 形状: (batch_size, seq_length)

    # 找到每个样本中 label 的位置（第一个非 -100 的位置）
    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:
            # 取第一个 label 位置
            pos = label_positions[0]

            # 获取该位置的预测 token
            pred_token = np.argmax(logits[i, pos - 1, :])  # 注意：预测下一个 token
            true_token = labels[i, pos]

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

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

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

    # 由于我们关心的是二分类（0 或 1），需要将 token id 映射回 0/1
    # 这里假设 tokenizer 会将 "0" 和 "1" 编码为特定的 token id
    # 我们直接比较预测和真实 token 是否相同

    # 计算各项指标
    accuracy = accuracy_score(references, predictions)

    # 对于 precision/recall/f1，我们需要将 token 映射为二分类标签
    # 这里简化处理：检查预测是否等于参考值
    correct = (predictions == references).astype(int)

    # 使用 average='binary' 假设 label=1 是正类
    # 但由于我们比较的是 token id，这里用 accuracy 作为主要指标
    precision = precision_score(references, predictions, average="macro", zero_division=0)
    recall = recall_score(references, predictions, average="macro", zero_division=0)
    f1 = f1_score(references, predictions, average="macro", zero_division=0)

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


def main():
    """主函数"""
    # 解析参数
    args = parse_args()

    # 加载默认配置
    model_config, data_config, training_config = get_default_configs()

    # 根据命令行参数更新配置
    model_config, data_config, training_config = update_configs_from_args(
        args, model_config, data_config, training_config
    )

    # 设置随机种子
    set_seed(training_config.seed)

    # 打印配置信息
    print("=" * 60)
    print("训练配置")
    print("=" * 60)
    print(f"模型路径: {model_config.model_name_or_path}")
    print(f"数据目录: {data_config.data_dir}")
    print(f"输出目录: {training_config.output_dir}")
    print(f"训练轮数: {training_config.num_train_epochs}")
    print(f"Batch Size: {training_config.per_device_train_batch_size}")
    print(f"梯度累积: {training_config.gradient_accumulation_steps}")
    print(f"学习率: {training_config.learning_rate}")
    print(f"最大序列长度: {data_config.max_seq_length}")
    print("=" * 60)

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

    # 加载 tokenizer
    print("\n加载 Tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        model_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("加载模型...")
    model = AutoModelForCausalLM.from_pretrained(
        model_config.model_name_or_path,
        trust_remote_code=True,
        torch_dtype=torch.float16 if training_config.fp16 else torch.bfloat16 if training_config.bf16 else torch.float32,
        device_map="auto",
    )

    # 启用梯度检查点以节省显存
    model.gradient_checkpointing_enable()

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

    # 加载数据集
    print("\n加载数据集...")
    train_dataset, valid_dataset, _ = load_datasets(tokenizer, data_config)

    # 创建数据整理器
    data_collator = IntentDataCollator(
        tokenizer=tokenizer,
        padding="longest",
        max_length=data_config.max_seq_length,
    )

    # 配置训练参数
    training_args = TrainingArguments(
        output_dir=training_config.output_dir,
        num_train_epochs=training_config.num_train_epochs,
        per_device_train_batch_size=training_config.per_device_train_batch_size,
        per_device_eval_batch_size=training_config.per_device_eval_batch_size,
        gradient_accumulation_steps=training_config.gradient_accumulation_steps,
        learning_rate=training_config.learning_rate,
        weight_decay=training_config.weight_decay,
        warmup_ratio=training_config.warmup_ratio,
        lr_scheduler_type=training_config.lr_scheduler_type,
        optim=training_config.optim,
        max_grad_norm=training_config.max_grad_norm,
        logging_steps=training_config.logging_steps,
        save_steps=training_config.save_steps,
        eval_steps=training_config.eval_steps,
        evaluation_strategy=training_config.evaluation_strategy,
        save_strategy=training_config.save_strategy,
        save_total_limit=training_config.save_total_limit,
        load_best_model_at_end=training_config.load_best_model_at_end,
        metric_for_best_model=training_config.metric_for_best_model,
        greater_is_better=training_config.greater_is_better,
        fp16=training_config.fp16,
        bf16=training_config.bf16,
        dataloader_num_workers=training_config.dataloader_num_workers,
        seed=training_config.seed,
        report_to=training_config.report_to,
        logging_dir=os.path.join(training_config.output_dir, "logs"),
        remove_unused_columns=False,
    )

    # 创建 Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    # 开始训练
    print("\n开始训练...")
    print("=" * 60)

    if args.resume_from_checkpoint:
        trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
    else:
        trainer.train()

    # 保存最终模型
    print("\n保存模型...")
    final_model_path = os.path.join(training_config.output_dir, "final_model")
    trainer.save_model(final_model_path)
    tokenizer.save_pretrained(final_model_path)

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

    # 在验证集上进行最终评估
    print("\n在验证集上进行最终评估...")
    eval_results = trainer.evaluate()
    print("评估结果:")
    for key, value in eval_results.items():
        print(f"  {key}: {value:.4f}")


if __name__ == "__main__":
    main()