src/deep_learning/data/wiki/dataset.py

"""Wiki 数据集主模块

实现 WikiDataset 类，继承自 TextGenerationDataSource。
"""

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

import tensorflow as tf

from deep_learning.data.spec import TextGenerationDataSource, TokenizerBundle
from .loader import doc_load
from .transformer import transform
from .tokenizer import sentence_piece, character_vectorization


@dataclass
class WikiDataset(TextGenerationDataSource):
    """Wiki 数据集

    将文档加载、分词、统计等功能绑定在一起的数据集类。

    Usage:
        dataset = WikiDataset(
            data_dir="~/data/wiki/mini_c4",
            tokenizer_type="sentence_piece"  # 或 "character"
        )

        # 获取文档数据集
        doc_ds = dataset.doc_ds()

        # 获取 token 数据集
        tokens_ds = dataset.tokens_ds()

        # 打印统计信息
        dataset.stat(seq_length=256)
    """

    data_dir: str
    sequence_length: int = 256
    batch_size: int = 32
    validation_batches: int = 0
    glob_pattern: str = "*"
    tokenizer_type: str = "sentence_piece"

    _data_path: pathlib.Path = field(init=False, repr=False)
    _tokenizer_bundle: Optional[TokenizerBundle] = field(
        init=False, repr=False, default=None
    )

    def __post_init__(self):
        self._data_path = pathlib.Path(self.data_dir).expanduser()

    def _load_tokenizer(self):
        """懒加载分词器"""
        if self._tokenizer_bundle is None:
            if self.tokenizer_type == "sentence_piece":
                tokenizer, end_of_text, decode = sentence_piece()
            elif self.tokenizer_type == "character":
                tokenizer, end_of_text, decode = character_vectorization()
            else:
                raise ValueError(f"Unknown tokenizer type: {self.tokenizer_type}")

            vocab_size = tokenizer.vocabulary_size()
            self._tokenizer_bundle = TokenizerBundle(
                tokenizer=tokenizer,
                decode=decode,
                end_of_text=end_of_text,
                vocab_size=vocab_size
            )

    def doc_ds(self) -> tf.data.Dataset:
        """返回原始文档数据集

        Returns:
            TensorFlow Dataset，每个元素是一个文档字符串
        """
        return doc_load(self._data_path, glob_pattern=self.glob_pattern)

    def tokens_ds(self) -> tf.data.Dataset:
        """返回 tokenized 数据集

        Returns:
            TensorFlow Dataset，每个元素是 (input_ids, target_ids) 对
        """
        self._load_tokenizer()
        ds = self.doc_ds()
        return transform(
            ds=ds,
            tokenizer=self._tokenizer_bundle.tokenizer,
            end_of_text=self._tokenizer_bundle.end_of_text,
            sequence_length=self.sequence_length,
            batch_size=self.batch_size
        )

    def tokenizer_bundle(self) -> TokenizerBundle:
        """返回分词器信息"""
        self._load_tokenizer()
        return self._tokenizer_bundle