src/korbyte/benchmark.py

"""Intrinsic tokenizer benchmarks with explicit baseline comparability."""

from __future__ import annotations

import gc
import importlib.metadata
import json
import math
import platform
import re
import statistics
import time
from collections.abc import Callable, Iterable
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Protocol

import psutil
from datasets import load_dataset
from tokenizers import Tokenizer

from .config import (
    DEFAULT_SEED,
    KANANA_MODEL_ID,
    KANANA_REVISION,
    KLUE_REVISION,
    Paths,
)

WHITESPACE_UNIT_RE = re.compile(r"\S+")

KLUE_FIELDS: dict[str, tuple[str, ...]] = {
    "ynat": ("title",),
    "sts": ("sentence1", "sentence2"),
    "nli": ("premise", "hypothesis"),
    "ner": ("sentence",),
    "re": ("sentence",),
    "dp": ("sentence",),
    "mrc": ("context", "question"),
    "wos": ("dialogue",),
}


class Counter(Protocol):
    name: str
    reversible: bool

    def count_many(self, texts: list[str]) -> list[int]: ...


@dataclass
class DomainMetrics:
    texts: int
    characters: int
    utf8_bytes: int
    whitespace_units: int
    tokens: int
    fertility: float
    characters_per_token: float
    bytes_per_token: float
    throughput_mib_s: float
    elapsed_s_median: float


@dataclass
class FastCounter:
    name: str
    tokenizer: Tokenizer
    reversible: bool = True

    def count_many(self, texts: list[str]) -> list[int]:
        return [len(encoding.ids) for encoding in self.tokenizer.encode_batch(texts)]


@dataclass
class FunctionCounter:
    name: str
    function: Callable[[str], list[str]]
    reversible: bool = False

    def count_many(self, texts: list[str]) -> list[int]:
        return [len(self.function(text)) for text in texts]


def _flatten_strings(value: Any) -> Iterable[str]:
    if isinstance(value, str):
        if value.strip():
            yield value
    elif isinstance(value, list):
        for item in value:
            yield from _flatten_strings(item)
    elif isinstance(value, dict):
        text = value.get("text")
        if isinstance(text, str):
            yield from _flatten_strings(text)


def load_klue_domains(*, limit_per_domain: int, seed: int = DEFAULT_SEED) -> dict[str, list[str]]:
    """Load deterministic held-out Korean text without persisting examples."""

    domains: dict[str, list[str]] = {}
    for index, (config_name, fields) in enumerate(KLUE_FIELDS.items()):
        dataset = load_dataset(
            "klue/klue",
            config_name,
            split="validation",
            revision=KLUE_REVISION,
        )
        dataset = dataset.shuffle(seed=seed + index)
        texts: list[str] = []
        for row in dataset:
            for field in fields:
                texts.extend(_flatten_strings(row.get(field)))
                if len(texts) >= limit_per_domain:
                    break
            if len(texts) >= limit_per_domain:
                break
        domains[config_name] = texts[:limit_per_domain]
        if not domains[config_name]:
            raise RuntimeError(f"No benchmark text extracted for KLUE/{config_name}")
    return domains


def _timed_counts(counter: Counter, texts: list[str], *, repeats: int) -> tuple[list[int], float]:
    warmup = texts[: min(32, len(texts))]
    if warmup:
        counter.count_many(warmup)
    durations: list[float] = []
    counts: list[int] | None = None
    for _ in range(repeats):
        gc.collect()
        started = time.perf_counter()
        current = counter.count_many(texts)
        durations.append(time.perf_counter() - started)
        if counts is None:
            counts = current
        elif counts != current:
            raise RuntimeError(f"Non-deterministic token counts from {counter.name}")
    if counts is None:
        return [], 0.0
    return counts, statistics.median(durations)


def _metrics(counter: Counter, texts: list[str], *, repeats: int) -> DomainMetrics:
    counts, elapsed = _timed_counts(counter, texts, repeats=repeats)
    characters = sum(map(len, texts))
    utf8_bytes = sum(len(text.encode("utf-8")) for text in texts)
    units = sum(len(WHITESPACE_UNIT_RE.findall(text)) for text in texts)
    tokens = sum(counts)
    mib = utf8_bytes / (1024 * 1024)
    return DomainMetrics(
        texts=len(texts),
        characters=characters,
        utf8_bytes=utf8_bytes,
        whitespace_units=units,
        tokens=tokens,
        fertility=tokens / units if units else math.nan,
        characters_per_token=characters / tokens if tokens else math.nan,
        bytes_per_token=utf8_bytes / tokens if tokens else math.nan,
        throughput_mib_s=mib / elapsed if elapsed else math.inf,
        elapsed_s_median=elapsed,
    )


def _optional_counters() -> tuple[list[Counter], dict[str, str]]:
    counters: list[Counter] = []
    unavailable: dict[str, str] = {}
    try:
        from konlpy.tag import Okt

        okt = Okt()
        counters.append(
            FunctionCounter("okt", lambda text: okt.morphs(text, norm=False, stem=False))
        )
    except Exception as error:  # pragma: no cover - host dependency
        unavailable["okt"] = f"{type(error).__name__}: {error}"

    try:
        try:
            import MeCab
        except ModuleNotFoundError:
            import mecab_ko as MeCab

        tagger = MeCab.Tagger("-Owakati")
        counters.append(
            FunctionCounter("mecab-ko", lambda text: tagger.parse(text).strip().split())
        )
    except Exception as error:  # pragma: no cover - host dependency
        unavailable["mecab-ko"] = f"{type(error).__name__}: {error}"
    return counters, unavailable


def _package_versions() -> dict[str, str]:
    versions: dict[str, str] = {}
    for package in ("tokenizers", "konlpy", "mecab-ko", "mecab-ko-dic"):
        try:
            versions[package] = importlib.metadata.version(package)
        except importlib.metadata.PackageNotFoundError:
            continue
    return versions


def run_benchmark(
    root: Path,
    *,
    limit_per_domain: int = 1_000,
    repeats: int = 3,
    include_morphological: bool = True,
) -> dict[str, Any]:
    """Run and persist the full held-out intrinsic benchmark."""

    paths = Paths(root)
    tokenizer_path = root / "tokenizer.json"
    if not tokenizer_path.is_file():
        raise FileNotFoundError(tokenizer_path)
    domains = load_klue_domains(limit_per_domain=limit_per_domain)
    counters: list[Counter] = [
        FastCounter("korbyte-128k", Tokenizer.from_file(str(tokenizer_path))),
        FastCounter(
            "kanana-2",
            Tokenizer.from_pretrained(KANANA_MODEL_ID, revision=KANANA_REVISION),
        ),
    ]
    unavailable: dict[str, str] = {}
    if include_morphological:
        optional, unavailable = _optional_counters()
        counters.extend(optional)

    results: dict[str, dict[str, dict[str, Any]]] = {}
    for counter in counters:
        results[counter.name] = {}
        for domain_name, texts in domains.items():
            metrics = _metrics(counter, texts, repeats=repeats)
            results[counter.name][domain_name] = asdict(metrics)
            print(
                f"{counter.name}/{domain_name}: {metrics.tokens:,} tokens, "
                f"{metrics.throughput_mib_s:.2f} MiB/s"
            )

    reductions = {
        domain: 100
        * (1 - results["korbyte-128k"][domain]["tokens"] / results["kanana-2"][domain]["tokens"])
        for domain in domains
    }
    macro_reduction = statistics.fmean(reductions.values())
    report = {
        "schema_version": 1,
        "created_at": datetime.now(UTC).isoformat(),
        "evaluation_dataset": {
            "id": "klue/klue",
            "revision": KLUE_REVISION,
            "split": "validation",
            "limit_per_domain": limit_per_domain,
            "text_examples_redistributed": False,
            "usage_note": (
                "KLUE text was excluded from tokenizer training and used only for "
                "intrinsic development evaluation."
            ),
        },
        "baseline": {
            "kanana_model_id": KANANA_MODEL_ID,
            "kanana_revision": KANANA_REVISION,
        },
        "repeats": repeats,
        "results": results,
        "korbyte_reduction_vs_kanana_percent": reductions,
        "korbyte_macro_reduction_vs_kanana_percent": macro_reduction,
        "compression_gate_percent": 5.0,
        "compression_gate_passed": macro_reduction >= 5.0,
        "unavailable_baselines": unavailable,
        "comparability_note": (
            "OKT and MeCab-ko are morphological analyzers, not reversible fixed-vocabulary "
            "LLM tokenizers. Kanana-2 is the primary like-for-like baseline."
        ),
        "environment": {
            "platform": platform.platform(),
            "processor": platform.processor(),
            "python": platform.python_version(),
            "logical_cpu_count": psutil.cpu_count(logical=True),
            "physical_memory_bytes": psutil.virtual_memory().total,
            "packages": _package_versions(),
        },
    }
    paths.benchmark_json.parent.mkdir(parents=True, exist_ok=True)
    paths.benchmark_json.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return report