config.py

"""
配置文件 - Qwen 意图识别模型训练配置
"""

from dataclasses import dataclass, field
from typing import Optional
import os


@dataclass
class ModelConfig:
    """模型配置"""
    # 模型路径（本地路径或 Hugging Face 模型名）
    model_name_or_path: str = "/Users/admin/Desktop/Qwen3/model"

    # 是否使用 flash attention（如果支持）
    use_flash_attention: bool = True

    # 数据类型
    torch_dtype: str = "auto"  # auto, float16, bfloat16, float32


@dataclass
class DataConfig:
    """数据配置"""
    # 数据集路径
    data_dir: str = "/Users/admin/Desktop/Qwen3/dataset"
    train_file: str = "train.jsonl"
    valid_file: str = "valid.jsonl"
    test_file: str = "test.jsonl"

    # 序列最大长度
    max_seq_length: int = 256

    # 数据加载
    preprocessing_num_workers: int = 4

    @property
    def train_path(self) -> str:
        return os.path.join(self.data_dir, self.train_file)

    @property
    def valid_path(self) -> str:
        return os.path.join(self.data_dir, self.valid_file)

    @property
    def test_path(self) -> str:
        return os.path.join(self.data_dir, self.test_file)


@dataclass
class TrainingConfig:
    """训练配置"""
    # 输出目录
    output_dir: str = "/Users/admin/Desktop/Qwen3/output"

    # 训练超参数
    num_train_epochs: int = 3
    per_device_train_batch_size: int = 8
    per_device_eval_batch_size: int = 16
    gradient_accumulation_steps: int = 4  # 等效 batch size = 8 * 4 = 32

    # 学习率相关
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    lr_scheduler_type: str = "cosine"

    # 优化器
    optim: str = "adamw_torch"

    # 梯度裁剪
    max_grad_norm: float = 1.0

    # 日志和保存
    logging_steps: int = 10
    save_steps: int = 100
    eval_steps: int = 100
    save_total_limit: int = 3

    # 评估策略
    evaluation_strategy: str = "steps"
    save_strategy: str = "steps"
    load_best_model_at_end: bool = True
    metric_for_best_model: str = "f1"
    greater_is_better: bool = True

    # 其他
    fp16: bool = True  # 混合精度训练
    bf16: bool = False  # 如果 GPU 支持 bfloat16，可以设为 True
    dataloader_num_workers: int = 4
    seed: int = 42

    # 报告
    report_to: str = "tensorboard"


# Prompt 模板
SYSTEM_PROMPT = """你是游戏生成意图识别器。
任务：判断用户是否请求生成互动小游戏。
输出：1（生成游戏）或 0（拒绝生成）
只输出数字，无需解释。"""


def build_prompt(text: str, label: Optional[int] = None) -> str:
    """
    构建完整的 prompt

    Args:
        text: 用户输入文本
        label: 标签（0 或 1），如果为 None 则不添加标签（用于推理）

    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


def build_prompt_for_inference(text: str) -> str:
    """
    构建用于推理的 prompt（不包含标签）

    Args:
        text: 用户输入文本

    Returns:
        用于推理的 prompt 字符串
    """
    return build_prompt(text, label=None)


# 标签映射
LABEL2ID = {"0": 0, "1": 1}
ID2LABEL = {0: "0", 1: "1"}


# 默认配置实例
def get_default_configs():
    """获取默认配置"""
    return ModelConfig(), DataConfig(), TrainingConfig()