# SillokBert-Scratch 프로젝트 1단계: 실록 특화 토크나이저 훈련
# -----------------------------------------------------------------
# BPE 알고리즘과 PUA/NFC 전처리 파이프라인을 사용하여
# Sillok 데이터에 최적화된 토크나이저를 훈련합니다.
# -----------------------------------------------------------------
import os
import re
from pathlib import Path
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.normalizers import NFC

def train_sillok_bpe_tokenizer():
    """
    BPE 알고리즘과 강화된 전처리 파이프라인을 사용하여 토크나이저를 훈련합니다.
    """
    # --- 경로 설정 ---
    corpus_file = "/home/work/baro/sillok25060103/preprocessed_corpus/train.txt"
    project_dir = Path("/home/work/baro/sillok/sillok_scratch_20250626")
    output_dir = project_dir / "sillok_tokenizer_bpe_preprocessed"
    tokenizer_file = output_dir / "tokenizer.json"

    print("--- 1. Sillok-Specific Tokenizer Training ---")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # --- PUA 변환 로직 ---
    # 실제 데이터에 맞는 완전한 매핑 테이블로 교체해야 합니다.
    pua_to_standard_map = {}
    
    pua_regex = None
    if pua_to_standard_map:
        pua_regex = re.compile("|".join(map(re.escape, pua_to_standard_map.keys())))
    
    def normalize_pua(text: str) -> str:
        if not pua_regex:
            return text
        return pua_regex.sub(lambda m: pua_to_standard_map[m.group(0)], text)

    # --- 토크나이저 설정 ---
    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
    tokenizer.normalizer = NFC()
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)

    trainer = BpeTrainer(
        vocab_size=500000,
        min_frequency=1,
        special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
    )

    # --- 코퍼스 이터레이터 및 훈련 ---
    def get_training_corpus():
        with open(corpus_file, "r", encoding="utf-8") as f:
            for line in f:
                processed_line = normalize_pua(line)
                yield processed_line.strip()

    print("BPE 토크나이저 훈련을 시작합니다...")
    tokenizer.train_from_iterator(get_training_corpus(), trainer=trainer)
    
    # --- 저장 및 검증 ---
    tokenizer.save(str(tokenizer_file))
    print(f"\n🎉 BPE 토크나이저 훈련 완료. 결과 파일: {tokenizer_file}")
    
    reloaded_tokenizer = Tokenizer.from_file(str(tokenizer_file))
    print(f"실제 생성된 어휘집 크기: {reloaded_tokenizer.get_vocab_size()}")

if __name__ == "__main__":
    train_sillok_bpe_tokenizer()