tokenizer/thai_tokenizer.py

import os
import json
from typing import List, Dict, Optional
import sentencepiece as spm
import torch
from .preprocess import preprocess_thai   # จาก Phase 1


# Special token IDs — ตรงกับที่ตั้งค่าตอน SentencePiece training
PAD_ID  = 0
UNK_ID  = 1
BOS_ID  = 2
EOS_ID  = 3
CLS_ID  = 4   # user_defined_symbols ลำดับที่ 1
SEP_ID  = 5   # user_defined_symbols ลำดับที่ 2
MASK_ID = 6   # user_defined_symbols ลำดับที่ 3


class ThaiTokenizer:
    """
    Wraps SentencePiece BPE สำหรับภาษาไทย
    - preprocess text ก่อน encode ทุกครั้ง
    - รองรับ special tokens: [PAD], [UNK], [BOS], [EOS], [CLS], [SEP], [MASK]
    - batch_encode คืน torch.Tensor พร้อมใช้ใน DataLoader
    """

    def __init__(self, model_path: str):
        """
        Parameters
        ----------
        model_path : str
            path ถึง .model file จาก SentencePiece training
            เช่น "tokenizer/thai_bpe.model"
        """
        if not os.path.exists(model_path):
            raise FileNotFoundError(
                f"ไม่พบ model file: {model_path}\n"
                f"รัน train_tokenizer.py ก่อนเพื่อสร้าง model"
            )

        self.model_path = model_path
        self.sp = spm.SentencePieceProcessor()
        self.sp.load(model_path)

        # validate vocab size
        self.vocab_size = self.sp.get_piece_size()

        # special token ids (ดึงจาก model จริง ไม่ hardcode)
        self.pad_id  = self.sp.piece_to_id("[PAD]")
        self.unk_id  = self.sp.piece_to_id("[UNK]")
        self.bos_id  = self.sp.piece_to_id("[BOS]")
        self.eos_id  = self.sp.piece_to_id("[EOS]")
        self.cls_id  = self.sp.piece_to_id("[CLS]")
        self.sep_id  = self.sp.piece_to_id("[SEP]")
        self.mask_id = self.sp.piece_to_id("[MASK]")

    # ──────────────────────────────────────────────
    # Core encode / decode
    # ──────────────────────────────────────────────

    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
        text = preprocess_thai(text)
        if not text:
            return [self.cls_id, self.sep_id] if add_special_tokens else []

        ids = self.sp.encode(text, out_type=int)

        if add_special_tokens:
            ids = [self.cls_id] + ids + [self.sep_id]  # ← เปลี่ยนตรงนี้

        return ids

    def encode_qa(
    self,
    question: str,
    context: str,
    max_length: int = 512,
    ) -> dict:
        """
        QA ต้องการ: [CLS] question [SEP] context [SEP]
        และต้องรู้ว่า context เริ่มที่ position ไหน
        """
        q_ids = self.sp.encode(preprocess_thai(question), out_type=int)
        c_ids = self.sp.encode(preprocess_thai(context),  out_type=int)

        # [CLS] q_ids [SEP] c_ids [SEP]
        ids = [self.cls_id] + q_ids + [self.sep_id] + c_ids + [self.sep_id]

        # context เริ่มที่ position นี้ (QAHead ใช้ mask คำนวณ span)
        context_start = 1 + len(q_ids) + 1   # หลัง [CLS] + question + [SEP]

        # truncate ถ้ายาวเกิน
        if len(ids) > max_length:
            # ตัด context ก่อน ไม่ตัด question
            allowed_ctx = max_length - context_start - 1  # -1 สำหรับ [SEP] สุดท้าย
            c_ids = c_ids[:allowed_ctx]
            ids = [self.cls_id] + q_ids + [self.sep_id] + c_ids + [self.sep_id]

        attention_mask = [1] * len(ids)

        return {
            "input_ids":      ids,
            "attention_mask": attention_mask,
            "context_start":  context_start,   # QAHead ใช้ป้องกัน predict span ใน question
        }

    def decode(
        self,
        ids: List[int],
        skip_special_tokens: bool = True,
    ) -> str:
        """
        แปลง list of token IDs → text

        Parameters
        ----------
        skip_special_tokens : bool
            ถ้า True จะกรอง PAD, BOS, EOS, CLS, SEP, MASK ออกก่อน decode
        """
        special = {
            self.pad_id, self.bos_id, self.eos_id,
            self.cls_id, self.sep_id, self.mask_id,
        }

        if skip_special_tokens:
            ids = [i for i in ids if i not in special]

        return self.sp.decode(ids)

    # ──────────────────────────────────────────────
    # Batch encode — ใช้ใน DataLoader
    # ──────────────────────────────────────────────

    def batch_encode(
        self,
        texts: List[str],
        max_length: int = 512,
        padding: bool = True,
        add_special_tokens: bool = True,
        return_tensors: bool = True,
    ) -> Dict:
        """
        Encode หลาย texts พร้อมกัน พร้อม padding และ truncation

        Returns
        -------
        dict with keys:
            input_ids      : (B, T) — token IDs
            attention_mask : (B, T) — 1=real token, 0=padding
        """
        all_ids = []
        for text in texts:
            ids = self.encode(text, add_special_tokens=add_special_tokens)

            # Truncate: ถ้ายาวเกิน max_length ให้ตัดท้าย
            # แต่รักษา EOS ไว้เสมอ
            if len(ids) > max_length:
                if add_special_tokens:
                    ids = ids[:max_length - 1] + [self.eos_id]
                else:
                    ids = ids[:max_length]

            all_ids.append(ids)

        # Padding: pad ให้ยาวเท่ากับ sequence ที่ยาวที่สุดใน batch
        if padding:
            max_len = max(len(ids) for ids in all_ids)
            padded_ids = []
            attention_masks = []

            for ids in all_ids:
                pad_len = max_len - len(ids)
                padded_ids.append(ids + [self.pad_id] * pad_len)
                # attention_mask: 1 สำหรับ real token, 0 สำหรับ padding
                attention_masks.append([1] * len(ids) + [0] * pad_len)
        else:
            padded_ids = all_ids
            attention_masks = [[1] * len(ids) for ids in all_ids]

        if return_tensors:
            return {
                "input_ids":      torch.tensor(padded_ids,      dtype=torch.long),
                "attention_mask": torch.tensor(attention_masks, dtype=torch.long),
            }

        return {
            "input_ids":      padded_ids,
            "attention_mask": attention_masks,
        }

    # ──────────────────────────────────────────────
    # Save / Load
    # ──────────────────────────────────────────────

    def save(self, output_dir: str) -> None:
        """
        บันทึก tokenizer ลง directory
        สร้าง 2 ไฟล์:
          - thai_bpe.model  (SentencePiece binary)
          - tokenizer_config.json  (metadata)
        """
        os.makedirs(output_dir, exist_ok=True)

        # copy model file
        import shutil
        dst = os.path.join(output_dir, "thai_bpe.model")
        if os.path.abspath(self.model_path) != os.path.abspath(dst):
            shutil.copy2(self.model_path, dst)

        # บันทึก metadata
        config = {
            "vocab_size": self.vocab_size,
            "model_file": "thai_bpe.model",
            "special_tokens": {
                "pad_id":  self.pad_id,
                "unk_id":  self.unk_id,
                "bos_id":  self.bos_id,
                "eos_id":  self.eos_id,
                "cls_id":  self.cls_id,
                "sep_id":  self.sep_id,
                "mask_id": self.mask_id,
            }
        }
        with open(os.path.join(output_dir, "tokenizer_config.json"), "w") as f:
            json.dump(config, f, indent=2, ensure_ascii=False)

        print(f"tokenizer saved → {output_dir}")

    @classmethod
    def from_pretrained(cls, directory: str) -> "ThaiTokenizer":
        """
        โหลด tokenizer จาก directory ที่ save() ไว้
        
        Usage:
            tok = ThaiTokenizer.from_pretrained("outputs/tokenizer")
        """
        config_path = os.path.join(directory, "tokenizer_config.json")
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"ไม่พบ tokenizer_config.json ใน {directory}")

        with open(config_path) as f:
            config = json.load(f)

        model_path = os.path.join(directory, config["model_file"])
        return cls(model_path=model_path)

    # ──────────────────────────────────────────────
    # Utilities
    # ──────────────────────────────────────────────

    def tokenize(self, text: str) -> List[str]:
        """แสดง subword pieces เป็น string — ใช้สำหรับ debug"""
        text = preprocess_thai(text)
        return self.sp.encode(text, out_type=str)

    def vocab_size_actual(self) -> int:
        return self.sp.get_piece_size()

    def __len__(self) -> int:
        return self.vocab_size

    def __repr__(self) -> str:
        return (
            f"ThaiTokenizer("
            f"vocab_size={self.vocab_size}, "
            f"model='{os.path.basename(self.model_path)}')"
        )