"""
RoBERTa-wwm-ext-large 商品分类微调脚本 (ModelScope版)
硬件配置: 8核CPU / 32GB内存 / 24GB显存
模型来源: ModelScope
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.utils.class_weight import compute_class_weight
from sklearn.metrics import accuracy_score, f1_score, classification_report
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
    set_seed,
    TrainerCallback
)
from transformers.modeling_outputs import SequenceClassifierOutput
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from tqdm import tqdm
import warnings
import re
import joblib
import os
import time
import glob
from datetime import datetime
from torch.utils.tensorboard import SummaryWriter
from safetensors.torch import save_file as safe_save_file

# ==================== ModelScope 导入 ====================
try:
    from modelscope.hub.snapshot_download import snapshot_download
    MODELSCOPE_AVAILABLE = True
    print("✅ ModelScope 已安装")
except ImportError:
    MODELSCOPE_AVAILABLE = False
    print("⚠️ ModelScope 未安装，将使用HuggingFace缓存")
    print("   安装命令: pip install modelscope")


# ==================== 1. 参数配置 ====================
class Config:
    # ===== 模型配置 =====
    MODEL_ID = "dienstag/chinese-roberta-wwm-ext"

    # 本地缓存目录
    CACHE_DIR = "./modelscope_cache"

    # 是否使用ModelScope下载
    USE_MODELSCOPE = True

    # ===== 训练参数 =====
    MAX_LENGTH = 64
    BATCH_SIZE = 32
    GRADIENT_ACCUMULATION_STEPS = 2
    NUM_EPOCHS = 5
    LEARNING_RATE = 2e-5
    WARMUP_RATIO = 0.1
    WEIGHT_DECAY = 0.01
    FP16 = True
    EARLY_STOPPING_PATIENCE = 3
    SAVE_TOTAL_LIMIT = 5

    # 日志和保存步数
    LOGGING_STEPS = 50
    EVAL_SAVE_STEPS = 200

    # 数据加载
    NUM_WORKERS = 6
    TEST_SIZE = 0.2

    # 路径
    DATA_FILE = "train_v4.xlsx"
    OUTPUT_DIR = "./results/roberta_large"
    LOG_DIR = "./logs/roberta_large"
    SAVE_DIR = "./saved_model/roberta_large"

    # 随机种子
    SEED = 42

    # 设备
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    # 评估指标配置
    METRIC_FOR_BEST_MODEL = "f1_macro"
    GREATER_IS_BETTER = True


# ==================== 2. ModelScope 模型加载 ====================
def download_model_from_modelscope(model_id, cache_dir):
    """从ModelScope下载模型"""
    if not MODELSCOPE_AVAILABLE:
        raise ImportError("请先安装 modelscope: pip install modelscope")

    print(f"📥 从ModelScope下载模型: {model_id}")
    print(f"📁 缓存目录: {cache_dir}")

    try:
        model_dir = snapshot_download(
            model_id,
            cache_dir=cache_dir,
            revision='master'
        )
        print(f"✅ 模型下载完成: {model_dir}")
        return model_dir
    except Exception as e:
        print(f"⚠️ ModelScope下载失败: {e}")
        print("   尝试使用HuggingFace缓存...")
        return None


def get_model_path(model_id, cache_dir, use_modelscope=True):
    """获取模型路径"""
    if use_modelscope and MODELSCOPE_AVAILABLE:
        try:
            local_dir = download_model_from_modelscope(model_id, cache_dir)
            if local_dir and os.path.exists(local_dir):
                return local_dir
        except Exception as e:
            print(f"⚠️ ModelScope加载失败: {e}")

    print(f"🔄 使用HuggingFace加载: {model_id}")
    return model_id


# ==================== 3. 工具函数 ====================
def print_gpu_memory():
    """打印GPU显存使用情况"""
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024 ** 3
        reserved = torch.cuda.memory_reserved() / 1024 ** 3
        max_allocated = torch.cuda.max_memory_allocated() / 1024 ** 3
        print(f"📊 GPU显存: 已分配 {allocated:.2f}GB | 已预留 {reserved:.2f}GB | 峰值 {max_allocated:.2f}GB")


def format_time(seconds):
    """格式化时间显示"""
    if seconds < 60:
        return f"{seconds:.0f}秒"
    elif seconds < 3600:
        return f"{seconds / 60:.1f}分钟"
    else:
        return f"{seconds / 3600:.2f}小时"


# ==================== 4. 数据加载 ====================
def load_data(file_path):
    """加载数据文件"""
    try:
        if file_path.endswith('.csv'):
            df = pd.read_csv(file_path)
        elif file_path.endswith(('.xlsx', '.xls')):
            df = pd.read_excel(file_path)
        else:
            raise ValueError("不支持的文件格式，请使用 .csv 或 .xlsx")

        print(f"📄 原始数据列: {df.columns.tolist()}")

        # 自动识别文本列
        if 'sentence' not in df.columns:
            possible_text_cols = ['sentence', 'content', 'name', 'product_name',
                                  'goods_name', 'title', 'Title', '商品名称', '商品全称']
            for col in possible_text_cols:
                if col in df.columns:
                    df.rename(columns={col: 'sentence'}, inplace=True)
                    break

        # 自动识别标签列
        if 'label' not in df.columns:
            possible_label_cols = ['label', 'class', 'type', 'cate',
                                   'Type', '分类', '类别']
            for col in possible_label_cols:
                if col in df.columns:
                    df.rename(columns={col: 'label'}, inplace=True)
                    break

        assert 'sentence' in df.columns and 'label' in df.columns, \
            "数据必须包含'sentence'(文本)和'label'(标签)列"

        df = df.dropna(subset=['sentence', 'label']).reset_index(drop=True)

        print(f"✅ 数据加载成功 | 样本量: {len(df)} | 分类数: {df['label'].nunique()}")
        return df
    except Exception as e:
        warnings.warn(f"❌ 数据加载失败: {str(e)}")
        raise


# ==================== 5. 数据清洗 ====================
def clean_chinese_text(text):
    """优化版清洗：保留中文字符、英文字母、数字、空格和重要符号"""
    if not isinstance(text, str):
        return ""

    cleaned = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\-/\.\(\)\（\）\·]', ' ', text)
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned


def clean_batch_texts(texts, show_sample=True):
    """批量清洗文本"""
    print("🧼 开始清洗文本数据...")
    cleaned = [clean_chinese_text(t) for t in texts]

    original_len = sum(len(t) for t in texts)
    cleaned_len = sum(len(t) for t in cleaned)
    if original_len > 0:
        print(f"   原始总字符: {original_len:,} | 清洗后总字符: {cleaned_len:,} | 缩减: {(1 - cleaned_len / original_len) * 100:.1f}%")

    if show_sample and len(texts) > 0:
        print("\n📝 清洗样例:")
        for i in range(min(3, len(texts))):
            print(f"  原始: {texts[i][:50]}{'...' if len(texts[i]) > 50 else ''}")
            print(f"  清洗: {cleaned[i][:50]}{'...' if len(cleaned[i]) > 50 else ''}")
            print()

    return cleaned


# ==================== 6. 数据集类 ====================
class TextDataset(Dataset):
    """动态编码数据集"""

    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            max_length=self.max_length,
            padding=False,
            truncation=True,
            return_tensors=None
        )
        return {
            "input_ids": encoding["input_ids"],
            "attention_mask": encoding["attention_mask"],
            "labels": self.labels[idx]
        }


# ==================== 7. Collate函数 ====================
def collate_fn(batch, tokenizer):
    """动态填充批处理"""
    input_ids = [item["input_ids"] for item in batch]
    attention_mask = [item["attention_mask"] for item in batch]
    labels = [item["labels"] for item in batch]

    padded = tokenizer.pad(
        {"input_ids": input_ids, "attention_mask": attention_mask},
        padding=True,
        return_tensors="pt"
    )

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


# ==================== 8. 评估指标 ====================
def compute_metrics(eval_pred):
    """计算多种评估指标"""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)

    return {
        "accuracy": accuracy_score(labels, predictions),
        "f1_macro": f1_score(labels, predictions, average="macro"),
        "f1_weighted": f1_score(labels, predictions, average="weighted"),
        "f1_micro": f1_score(labels, predictions, average="micro"),
    }


# ==================== 9. 类别权重 ====================
def compute_class_weights(labels, num_classes):
    """计算类别权重"""
    classes = np.unique(labels)
    weights = compute_class_weight("balanced", classes=classes, y=labels)

    weight_tensor = torch.ones(num_classes, dtype=torch.float)
    for cls, weight in zip(classes, weights):
        weight_tensor[cls] = weight

    return weight_tensor


# ==================== 10. 带权重损失的自定义模型 ====================
class CustomClassificationModel(torch.nn.Module):
    """包装原始模型，支持自定义损失函数"""

    def __init__(self, base_model, class_weights=None):
        super().__init__()
        self.base_model = base_model
        self.num_labels = base_model.config.num_labels
        self.config = base_model.config

        if class_weights is not None:
            self.loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights)
        else:
            self.loss_fct = torch.nn.CrossEntropyLoss()

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            **kwargs
        )

        logits = outputs.logits

        loss = None
        if labels is not None:
            loss = self.loss_fct(logits, labels)

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states if hasattr(outputs, 'hidden_states') else None,
            attentions=outputs.attentions if hasattr(outputs, 'attentions') else None
        )


# ==================== 11. 自定义保存回调（修复非连续张量问题） ====================
class CustomSaveCallback(TrainerCallback):
    """
    自定义保存回调：在保存前修复非连续张量
    """
    def on_save(self, args, state, control, model, tokenizer=None, **kwargs):
        if state.is_world_process_zero:
            # 获取当前保存路径
            checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")

            # 如果是包装模型，获取基础模型
            if hasattr(model, 'base_model'):
                model_to_save = model.base_model
            else:
                model_to_save = model

            # 确保所有参数连续
            for param in model_to_save.parameters():
                if not param.is_contiguous():
                    param.data = param.data.contiguous()

            # 保存模型（使用pickle格式）
            try:
                model_to_save.save_pretrained(checkpoint_dir)
                if tokenizer:
                    tokenizer.save_pretrained(checkpoint_dir)
                print(f"✅ 检查点已保存: {checkpoint_dir}")
            except Exception as e:
                print(f"⚠️ 保存检查点失败: {e}")
                # 备用方案：手动保存
                self._manual_save(model_to_save, tokenizer, checkpoint_dir)

    def _manual_save(self, model, tokenizer, save_dir):
        """手动保存模型"""
        try:
            os.makedirs(save_dir, exist_ok=True)
            model.config.save_pretrained(save_dir)

            state_dict = model.state_dict()
            contiguous_state_dict = {}
            for key, tensor in state_dict.items():
                if not tensor.is_contiguous():
                    contiguous_state_dict[key] = tensor.contiguous()
                else:
                    contiguous_state_dict[key] = tensor

            torch.save(contiguous_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
            if tokenizer:
                tokenizer.save_pretrained(save_dir)
            print(f"✅ 检查点已保存 (手动): {save_dir}")
        except Exception as e:
            print(f"❌ 保存失败: {e}")


# ==================== 12. TensorBoard回调 ====================
class TensorBoardCallback(TrainerCallback):
    """TensorBoard回调"""

    def __init__(self, log_dir):
        self.writer = SummaryWriter(log_dir)

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs:
            step = state.global_step
            for key, value in logs.items():
                if isinstance(value, (int, float)):
                    self.writer.add_scalar(key, value, step)

    def on_train_end(self, args, state, control, **kwargs):
        self.writer.close()


# ==================== 13. 模型初始化 ====================
def init_model(num_labels, class_weights=None):
    """初始化模型和分词器"""
    model_path = get_model_path(
        Config.MODEL_ID,
        Config.CACHE_DIR,
        Config.USE_MODELSCOPE
    )

    print(f"\n🔄 加载模型: {model_path}")
    start_time = time.time()

    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        trust_remote_code=True
    )

    base_model = AutoModelForSequenceClassification.from_pretrained(
        model_path,
        num_labels=num_labels,
        trust_remote_code=True
    )

    # 确保模型参数是连续的（修复非连续张量问题）
    base_model = base_model.to(Config.DEVICE)

    # 遍历所有参数，确保连续
    for param in base_model.parameters():
        if not param.is_contiguous():
            param.data = param.data.contiguous()

    if class_weights is not None:
        class_weights = class_weights.to(Config.DEVICE)
        model = CustomClassificationModel(base_model, class_weights)
    else:
        model = base_model

    load_time = time.time() - start_time
    print(f"✅ 模型加载完成，耗时: {format_time(load_time)}")

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"📊 模型参数总量: {total_params:,} ({total_params / 1e6:.1f}M)")
    print(f"📊 可训练参数: {trainable_params:,} ({trainable_params / 1e6:.1f}M)")

    return tokenizer, model


# ==================== 14. 训练配置 ====================
def get_training_args():
    """获取训练参数配置"""
    return TrainingArguments(
        output_dir=Config.OUTPUT_DIR,
        num_train_epochs=Config.NUM_EPOCHS,
        per_device_train_batch_size=Config.BATCH_SIZE,
        per_device_eval_batch_size=Config.BATCH_SIZE * 2,
        gradient_accumulation_steps=Config.GRADIENT_ACCUMULATION_STEPS,
        learning_rate=Config.LEARNING_RATE,
        warmup_ratio=Config.WARMUP_RATIO,
        weight_decay=Config.WEIGHT_DECAY,

        # 日志和保存
        logging_dir=Config.LOG_DIR,
        logging_steps=Config.LOGGING_STEPS,
        logging_first_step=True,
        eval_strategy="steps",
        eval_steps=Config.EVAL_SAVE_STEPS,
        save_strategy="steps",
        save_steps=Config.EVAL_SAVE_STEPS,
        save_total_limit=Config.SAVE_TOTAL_LIMIT,
        load_best_model_at_end=True,
        metric_for_best_model=Config.METRIC_FOR_BEST_MODEL,
        greater_is_better=Config.GREATER_IS_BETTER,

        # 性能优化
        fp16=Config.FP16,
        dataloader_num_workers=Config.NUM_WORKERS,
        dataloader_pin_memory=True,

        # 其他
        report_to="tensorboard",
        seed=Config.SEED,
        push_to_hub=False,
        lr_scheduler_type="linear",
        remove_unused_columns=False,
        dataloader_drop_last=False,

        # ✅ 关键修复：禁用 safetensors
        save_safetensors=False,
    )


# ==================== 15. 保存模型 ====================
def save_model(model, tokenizer, save_dir):
    """保存模型"""
    print(f"\n💾 保存模型到: {save_dir}")
    os.makedirs(save_dir, exist_ok=True)

    if hasattr(model, 'base_model'):
        model_to_save = model.base_model
    else:
        model_to_save = model

    # 确保所有参数连续
    for param in model_to_save.parameters():
        if not param.is_contiguous():
            param.data = param.data.contiguous()

    try:
        model_to_save.save_pretrained(save_dir)
        tokenizer.save_pretrained(save_dir)
        print(f"✅ 模型已保存到: {save_dir}")
        return True
    except Exception as e:
        print(f"⚠️ 保存失败: {e}")
        print("   尝试使用手动方式保存...")

    # 备用方案：手动保存
    try:
        model_to_save.config.save_pretrained(save_dir)

        state_dict = model_to_save.state_dict()
        contiguous_state_dict = {}
        for key, tensor in state_dict.items():
            if not tensor.is_contiguous():
                contiguous_state_dict[key] = tensor.contiguous()
            else:
                contiguous_state_dict[key] = tensor

        torch.save(contiguous_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
        tokenizer.save_pretrained(save_dir)
        print(f"✅ 模型已保存到: {save_dir}")
        return True
    except Exception as e:
        print(f"❌ 保存失败: {e}")
        return False


# ==================== 16. 预测函数 ====================
@torch.no_grad()
def batch_predict(texts, model, tokenizer, label_map, top_k=3, batch_size=64):
    """批量预测"""
    model.eval()
    all_results = []

    cleaned_texts = [clean_chinese_text(t) for t in texts]

    for i in tqdm(range(0, len(cleaned_texts), batch_size), desc="预测中"):
        batch = cleaned_texts[i:i + batch_size]
        inputs = tokenizer(
            batch,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=Config.MAX_LENGTH
        ).to(Config.DEVICE)

        outputs = model(**inputs)

        if hasattr(outputs, 'logits'):
            logits = outputs.logits
        else:
            logits = outputs

        probs = torch.softmax(logits, dim=1).cpu()

        for prob in probs:
            top_probs, top_indices = torch.topk(prob, k=min(top_k, len(label_map)))
            results = [
                {
                    "category": label_map[idx.item()],
                    "confidence": prob.item()
                }
                for prob, idx in zip(top_probs, top_indices)
            ]
            all_results.append(results)

    return all_results


# ==================== 17. 模型验证 ====================
def validate_model(model, tokenizer, sample_texts):
    """快速验证模型"""
    print("\n🔍 模型验证...")
    try:
        inputs = tokenizer(
            sample_texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=Config.MAX_LENGTH
        ).to(Config.DEVICE)

        outputs = model(**inputs)

        if hasattr(outputs, 'logits'):
            logits = outputs.logits
        else:
            logits = outputs

        print(f"✅ 模型验证通过，输出形状: {logits.shape}")
        return True
    except Exception as e:
        print(f"❌ 模型验证失败: {e}")
        return False


# ==================== 18. 数据统计 ====================
def print_data_stats(df):
    """打印数据统计信息"""
    print("\n" + "=" * 50)
    print("📊 数据统计")
    print("=" * 50)
    print(f"总样本量: {len(df)}")
    print(f"分类数: {df['label'].nunique()}")

    label_counts = df['label'].value_counts()
    print(f"\n📈 类别分布 (Top 10):")
    for label, count in label_counts.head(10).items():
        pct = count / len(df) * 100
        bar = "█" * int(pct / 2)
        print(f"  {label}: {count} ({pct:.1f}%) {bar}")

    tail_threshold = 10
    tail_count = (label_counts <= tail_threshold).sum()
    tail_samples = label_counts[label_counts <= tail_threshold].sum()
    print(f"\n📉 长尾类别 (≤{tail_threshold}条): {tail_count}个类别, 共{tail_samples}条样本")

    text_lengths = df['sentence'].str.len()
    print(f"\n📏 文本长度统计:")
    print(f"  平均: {text_lengths.mean():.1f}")
    print(f"  中位数: {text_lengths.median():.0f}")
    print(f"  最大: {text_lengths.max()}")
    print(f"  最小: {text_lengths.min()}")
    print("=" * 50 + "\n")


# ==================== 19. 主函数 ====================
def main():
    """主训练流程"""
    print("\n" + "=" * 60)
    print("🚀 RoBERTa-wwm-ext-large 商品分类微调 (ModelScope版)")
    print(f"📅 启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"💻 设备: {Config.DEVICE}")
    print(f"📦 模型来源: {'ModelScope' if Config.USE_MODELSCOPE else 'HuggingFace'}")
    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:.1f}GB")
    print("=" * 60)

    total_start = time.time()

    set_seed(Config.SEED)

    os.makedirs(Config.OUTPUT_DIR, exist_ok=True)
    os.makedirs(Config.LOG_DIR, exist_ok=True)
    os.makedirs(Config.SAVE_DIR, exist_ok=True)
    os.makedirs(Config.CACHE_DIR, exist_ok=True)

    # ===== 1. 加载数据 =====
    print(f"\n📂 加载数据: {Config.DATA_FILE}")
    df = load_data(Config.DATA_FILE)

    # ===== 2. 数据清洗 =====
    df['sentence_cleaned'] = clean_batch_texts(df['sentence'].tolist())
    df = df[df['sentence_cleaned'].str.len() > 0].reset_index(drop=True)
    print(f"✅ 清洗后剩余: {len(df)} 条")

    # ===== 3. 数据统计 =====
    print_data_stats(df)

    # ===== 4. 标签编码 =====
    print("🏷️ 处理标签...")
    label_encoder = LabelEncoder()
    df['label_id'] = label_encoder.fit_transform(df['label'])
    label_map = {i: label for i, label in enumerate(label_encoder.classes_)}
    num_labels = len(label_map)

    print(f"✅ 类别数: {num_labels}")
    print(f"   标签示例: {list(label_map.items())[:5]}")

    joblib.dump(label_encoder, "label_encoder_roberta_large.pkl")
    print("✅ 标签编码器已保存: label_encoder_roberta_large.pkl")

    # ===== 5. 类别权重 =====
    print("\n⚖️ 计算类别权重...")
    class_weights = compute_class_weights(df['label_id'].values, num_labels)
    print(f"   权重范围: {class_weights.min():.3f} - {class_weights.max():.3f}")
    print(f"   平均权重: {class_weights.mean():.3f}")

    # ===== 6. 划分数据集 =====
    print("\n📊 划分数据集...")
    train_df, eval_df = train_test_split(
        df,
        test_size=Config.TEST_SIZE,
        random_state=Config.SEED,
        stratify=df["label_id"]
    )
    print(f"   训练集: {len(train_df)} 条")
    print(f"   验证集: {len(eval_df)} 条")

    # ===== 7. 初始化模型 =====
    tokenizer, model = init_model(num_labels, class_weights)
    model.to(Config.DEVICE)
    print(f"✅ 模型已加载到: {Config.DEVICE}")
    print_gpu_memory()

    # ===== 8. 模型验证 =====
    sample_texts = ["九阳电饭煲家用", "苹果手机iPhone", "耐克运动鞋"]
    if not validate_model(model, tokenizer, sample_texts):
        print("⚠️ 模型验证失败，继续训练...")

    # ===== 9. 准备数据集 =====
    print("\n📦 准备数据集...")
    train_dataset = TextDataset(
        train_df['sentence_cleaned'].tolist(),
        train_df['label_id'].tolist(),
        tokenizer,
        Config.MAX_LENGTH
    )
    eval_dataset = TextDataset(
        eval_df['sentence_cleaned'].tolist(),
        eval_df['label_id'].tolist(),
        tokenizer,
        Config.MAX_LENGTH
    )

    # ===== 10. 配置训练器 =====
    print("\n⚙️ 配置训练器...")
    training_args = get_training_args()

    def data_collator(batch):
        return collate_fn(batch, tokenizer)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[
            EarlyStoppingCallback(early_stopping_patience=Config.EARLY_STOPPING_PATIENCE),
            TensorBoardCallback(Config.LOG_DIR),
            CustomSaveCallback(),  # ✅ 添加自定义保存回调
        ]
    )

    # ===== 11. 断点续训 =====
    checkpoint_dir = None
    checkpoints = glob.glob(os.path.join(Config.OUTPUT_DIR, "checkpoint-*"))
    if checkpoints:
        latest_checkpoint = max(checkpoints, key=os.path.getctime)
        print(f"📂 发现检查点: {latest_checkpoint}")
        checkpoint_dir = latest_checkpoint

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

    train_start = time.time()

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

    train_time = time.time() - train_start
    print(f"\n✅ 训练完成，耗时: {format_time(train_time)}")

    # ===== 13. 保存模型 =====
    save_model(model, tokenizer, Config.SAVE_DIR)

    # ===== 14. 最终评估 =====
    print("\n📊 最终评估...")
    eval_results = trainer.evaluate()
    print(f"   Eval Loss: {eval_results['eval_loss']:.4f}")
    print(f"   Eval Accuracy: {eval_results.get('eval_accuracy', 0):.4f}")
    print(f"   Eval F1 Macro: {eval_results.get('eval_f1_macro', 0):.4f}")
    print(f"   Eval F1 Weighted: {eval_results.get('eval_f1_weighted', 0):.4f}")

    # ===== 15. 推理测试 =====
    print("\n🧪 推理测试...")
    test_samples = [
        "九阳3升家用多功能电饭煲F30S-S160",
        "苹果iPhone 15 Pro Max 512GB 黑色钛金属",
        "海尔双开门冰箱BCD-600WGHSS19B8U1",
        "Lining李宁女款长裤 AKYV092-1 黑色",
        "珀莱雅双抗精华美白特证版补水保湿套装50ml礼盒",
    ]

    predictions = batch_predict(test_samples, model, tokenizer, label_map, top_k=3)

    print("\n📋 预测结果:")
    print("-" * 80)
    for sample, preds in zip(test_samples, predictions):
        print(f"输入: {sample}")
        for i, p in enumerate(preds):
            print(f"  Top{i + 1}: {p['category']} (置信度: {p['confidence']:.4f})")
        print("-" * 80)

    # ===== 16. 总结 =====
    total_time = time.time() - total_start
    print("\n" + "=" * 60)
    print("🎉 训练完成!")
    print(f"📅 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"⏱️ 总耗时: {format_time(total_time)}")
    print(f"⏱️ 训练耗时: {format_time(train_time)}")
    print(f"📊 最佳模型指标: {Config.METRIC_FOR_BEST_MODEL} = {trainer.state.best_metric:.4f}")
    print("=" * 60)


# ==================== 20. 入口 ====================
if __name__ == "__main__":
    main()