data_utils.py

"""
数据处理工具 - 加载和预处理意图识别数据集
"""

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

import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer

from config import build_prompt, DataConfig


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

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

        Args:
            data_path: JSONL 数据文件路径
            tokenizer: 分词器
            max_seq_length: 最大序列长度
        """
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        self.data = self._load_data(data_path)

    def _load_data(self, data_path: str) -> List[Dict[str, Any]]:
        """加载 JSONL 数据"""
        data = []
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    item = json.loads(line)
                    data.append(item)
        return data

    def __len__(self) -> int:
        return len(self.data)

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

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

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

        # 分词
        encoding = self.tokenizer(
            full_prompt,
            max_length=self.max_seq_length,
            padding=False,  # 在 collator 中进行 padding
            truncation=True,
            return_tensors=None,
        )

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

        # 创建 labels：只计算 label token 的 loss
        # 将 prompt 部分的 labels 设为 -100（忽略）
        labels = self._create_labels(text, label, 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, label: int, input_ids: List[int]
    ) -> List[int]:
        """
        创建训练标签

        只在 label token 位置计算 loss，其他位置设为 -100

        Args:
            text: 原始文本
            label: 标签（0 或 1）
            input_ids: 完整序列的 token ids

        Returns:
            labels 列表，非训练位置为 -100
        """
        # 构建不含 label 的 prompt
        prompt_without_label = build_prompt(text, label=None)

        # 对 prompt 进行分词，获取 prompt 的长度
        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"])

        # 创建 labels：prompt 部分为 -100，label 部分保留原始 token
        labels = [-100] * len(input_ids)

        # 从 prompt_length 开始的 token 用于计算 loss
        for i in range(prompt_length, len(input_ids)):
            labels[i] = input_ids[i]

        return labels


@dataclass
class IntentDataCollator:
    """
    数据整理器 - 处理动态 padding
    """

    tokenizer: PreTrainedTokenizer
    padding: str = "longest"  # longest, max_length
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = 8  # 优化 Tensor Core 计算

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

        Args:
            features: 样本列表

        Returns:
            整理后的 batch 字典
        """
        # 获取 batch 中的最大长度
        max_length_in_batch = max(len(f["input_ids"]) for f in features)

        # 对齐到 pad_to_multiple_of
        if self.pad_to_multiple_of is not None:
            max_length_in_batch = (
                (max_length_in_batch + self.pad_to_multiple_of - 1)
                // self.pad_to_multiple_of
                * self.pad_to_multiple_of
            )

        # 如果指定了 max_length，使用较小的值
        if self.max_length is not None:
            max_length_in_batch = min(max_length_in_batch, 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()

            # 计算需要 padding 的长度
            padding_length = max_length_in_batch - len(input_ids)

            # 执行 padding（右 padding）
            input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_length
            attention_mask = attention_mask + [0] * padding_length
            labels = labels + [-100] * padding_length  # padding 位置不计算 loss

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

        # 转换为 tensor
        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 load_datasets(
    tokenizer: PreTrainedTokenizer,
    data_config: DataConfig,
) -> tuple:
    """
    加载训练集、验证集和测试集

    Args:
        tokenizer: 分词器
        data_config: 数据配置

    Returns:
        (train_dataset, valid_dataset, test_dataset)
    """
    train_dataset = IntentDataset(
        data_path=data_config.train_path,
        tokenizer=tokenizer,
        max_seq_length=data_config.max_seq_length,
    )

    valid_dataset = IntentDataset(
        data_path=data_config.valid_path,
        tokenizer=tokenizer,
        max_seq_length=data_config.max_seq_length,
    )

    test_dataset = IntentDataset(
        data_path=data_config.test_path,
        tokenizer=tokenizer,
        max_seq_length=data_config.max_seq_length,
    )

    print(f"训练集样本数: {len(train_dataset)}")
    print(f"验证集样本数: {len(valid_dataset)}")
    print(f"测试集样本数: {len(test_dataset)}")

    return train_dataset, valid_dataset, test_dataset