src/korbyte/comparison.py

"""Pinned public comparison and post-selection audit benchmarks."""

from __future__ import annotations

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

import psutil
from huggingface_hub import HfApi, hf_hub_download
from rapidfuzz.distance.Levenshtein import distance
from transformers import AutoTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizerFast

from .config import (
    COMPARISON_BASELINES,
    KMMLU_ID,
    KMMLU_REVISION,
    PUBLIC_BENCHMARK_ID,
    PUBLIC_BENCHMARK_REVISION,
    Paths,
)

PUBLIC_DOMAINS = (
    "finance",
    "legal",
    "lyrics",
    "news",
    "social",
    "subtitles",
    "web",
    "wiki",
)
WHITESPACE_UNIT_RE = re.compile(r"\S+")


@dataclass(frozen=True)
class Document:
    text: str
    words: int


@dataclass
class ComparisonMetrics:
    documents: int
    words: int
    characters: int
    utf8_bytes: int
    tokens: int
    fertility: float
    characters_per_token: float
    bytes_per_token: float
    exact_document_ratio: float
    byte_fidelity: float
    character_fidelity: float
    unknown_token_ratio: float
    observed_vocabulary: int
    ebpb: float
    throughput_mib_s: float


def _public_documents() -> list[Document]:
    documents: list[Document] = []
    for domain in PUBLIC_DOMAINS:
        path = hf_hub_download(
            PUBLIC_BENCHMARK_ID,
            f"ko/{domain}.jsonl.gz",
            repo_type="dataset",
            revision=PUBLIC_BENCHMARK_REVISION,
        )
        with gzip.open(path, "rt", encoding="utf-8") as handle:
            for line in handle:
                row = json.loads(line)
                text = row["text"]
                if len(text) != int(row["char_count"]):
                    raise RuntimeError(f"Character-count mismatch in public {domain} data")
                if len(text.encode()) != int(row["byte_count"]):
                    raise RuntimeError(f"Byte-count mismatch in public {domain} data")
                documents.append(Document(text=text, words=int(row["word_count"])))
    return documents


def _kmmlu_test_files(api: HfApi | None = None) -> list[str]:
    client = api or HfApi()
    files = client.list_repo_files(KMMLU_ID, repo_type="dataset", revision=KMMLU_REVISION)
    selected = sorted(
        name for name in files if name.startswith("data/") and name.endswith("-test.csv")
    )
    if len(selected) != 45:
        raise RuntimeError(f"Expected 45 KMMLU test files, found {len(selected)}")
    return selected


def _kmmlu_documents() -> list[Document]:
    documents: list[Document] = []
    for filename in _kmmlu_test_files():
        path = hf_hub_download(
            KMMLU_ID,
            filename,
            repo_type="dataset",
            revision=KMMLU_REVISION,
        )
        with Path(path).open(encoding="utf-8", newline="") as handle:
            for row in csv.DictReader(handle):
                text = "\n".join(
                    [row["question"], *(f"{label}. {row[label]}" for label in "ABCD")]
                )
                documents.append(Document(text=text, words=len(WHITESPACE_UNIT_RE.findall(text))))
    if len(documents) != 35_030:
        raise RuntimeError(f"Expected 35,030 KMMLU test questions, found {len(documents):,}")
    return documents


def _batches(values: list[Document], size: int = 512) -> Iterable[list[Document]]:
    for start in range(0, len(values), size):
        yield values[start : start + size]


def _encode(tokenizer: PreTrainedTokenizerBase, texts: list[str]) -> list[list[int]]:
    return tokenizer(
        texts,
        add_special_tokens=False,
        padding=False,
        truncation=False,
        return_attention_mask=False,
        return_token_type_ids=False,
        verbose=False,
    )["input_ids"]


def evaluate_tokenizer(
    tokenizer: PreTrainedTokenizerBase,
    documents: list[Document],
    *,
    repeats: int = 3,
) -> ComparisonMetrics:
    """Evaluate one tokenizer with the public leaderboard's core intrinsic metrics."""

    if repeats < 1:
        raise ValueError("repeats must be positive")
    token_count = 0
    exact_documents = 0
    byte_edits = 0
    character_edits = 0
    unknown_tokens = 0
    observed: set[int] = set()
    durations = [0.0] * repeats
    unk_id = tokenizer.unk_token_id

    for batch in _batches(documents):
        texts = [document.text for document in batch]
        reference: list[list[int]] | None = None
        for repeat in range(repeats):
            started = time.perf_counter()
            current = _encode(tokenizer, texts)
            durations[repeat] += time.perf_counter() - started
            if reference is None:
                reference = current
            elif current != reference:
                raise RuntimeError("Tokenizer produced non-deterministic IDs")
        assert reference is not None
        decoded = tokenizer.batch_decode(
            reference,
            skip_special_tokens=False,
            clean_up_tokenization_spaces=False,
        )
        for document, token_ids, reconstructed in zip(batch, reference, decoded, strict=True):
            token_count += len(token_ids)
            observed.update(token_ids)
            if unk_id is not None:
                unknown_tokens += token_ids.count(unk_id)
            if reconstructed == document.text:
                exact_documents += 1
            else:
                byte_edits += distance(document.text.encode(), reconstructed.encode())
                character_edits += distance(document.text, reconstructed)

    words = sum(document.words for document in documents)
    characters = sum(len(document.text) for document in documents)
    utf8_bytes = sum(len(document.text.encode()) for document in documents)
    elapsed = statistics.median(durations)
    observed_size = max(len(observed), 2)
    return ComparisonMetrics(
        documents=len(documents),
        words=words,
        characters=characters,
        utf8_bytes=utf8_bytes,
        tokens=token_count,
        fertility=token_count / words,
        characters_per_token=characters / token_count,
        bytes_per_token=utf8_bytes / token_count,
        exact_document_ratio=exact_documents / len(documents),
        byte_fidelity=1 - byte_edits / utf8_bytes,
        character_fidelity=1 - character_edits / characters,
        unknown_token_ratio=unknown_tokens / token_count,
        observed_vocabulary=len(observed),
        ebpb=(token_count * math.log2(observed_size) + 8 * byte_edits) / utf8_bytes,
        throughput_mib_s=utf8_bytes / (1024 * 1024) / elapsed,
    )


def _ranks(results: dict[str, dict[str, Any]], metric: str) -> dict[str, int]:
    successful = [
        (key, float(value["metrics"][metric]))
        for key, value in results.items()
        if "metrics" in value
    ]
    reverse = metric in {"bytes_per_token", "byte_fidelity", "character_fidelity"}
    successful.sort(key=lambda item: item[1], reverse=reverse)
    return {key: rank for rank, (key, _) in enumerate(successful, start=1)}


def _evaluate_set(
    root: Path,
    documents: list[Document],
    *,
    repeats: int,
) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, int]]]:
    systems: list[tuple[str, str, str | None, PreTrainedTokenizerBase | None, str | None]] = []
    local = PreTrainedTokenizerFast.from_pretrained(str(root), local_files_only=True)
    systems.append(("korbyte-128k", "dawncr0w/KorByte-128K", None, local, None))
    for baseline in COMPARISON_BASELINES:
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                baseline.model_id,
                revision=baseline.revision,
                use_fast=True,
                trust_remote_code=False,
            )
            systems.append(
                (baseline.key, baseline.model_id, baseline.revision, tokenizer, None)
            )
        except Exception as error:
            systems.append(
                (
                    baseline.key,
                    baseline.model_id,
                    baseline.revision,
                    None,
                    f"{type(error).__name__}: {error}",
                )
            )

    results: dict[str, dict[str, Any]] = {}
    for key, model_id, revision, tokenizer, error in systems:
        if tokenizer is None:
            results[key] = {"model_id": model_id, "revision": revision, "error": error}
            continue
        try:
            metrics = evaluate_tokenizer(tokenizer, documents, repeats=repeats)
        except Exception as evaluation_error:
            if key == "korbyte-128k":
                raise
            results[key] = {
                "model_id": model_id,
                "revision": revision,
                "error": f"{type(evaluation_error).__name__}: {evaluation_error}",
            }
            continue
        results[key] = {
            "model_id": model_id,
            "revision": revision,
            "vocabulary_size": len(tokenizer),
            "metrics": asdict(metrics),
        }
        print(
            f"{key}: fertility={metrics.fertility:.4f}, ebpb={metrics.ebpb:.4f}, "
            f"exact={metrics.exact_document_ratio:.4%}"
        )

    ranks = {metric: _ranks(results, metric) for metric in ("fertility", "ebpb")}
    for metric, values in ranks.items():
        for key, rank in values.items():
            results[key].setdefault("ranks", {})[metric] = rank
    return results, ranks


def run_comparison(root: Path, *, repeats: int = 3) -> dict[str, Any]:
    """Run development-set and frozen post-selection audit comparisons."""

    public_results, public_ranks = _evaluate_set(root, _public_documents(), repeats=repeats)
    audit_results, audit_ranks = _evaluate_set(root, _kmmlu_documents(), repeats=repeats)
    ours_public = public_results["korbyte-128k"]["metrics"]
    ours_audit = audit_results["korbyte-128k"]["metrics"]
    required_systems = {
        "korbyte-128k",
        "ax-4.0",
        "trillion-7b",
        "midm-2.0",
        "k-exaone",
        "solar-pro3",
        "kanana-2",
        "hyperclovax",
    }
    coverage_passed = all(
        required_systems.issubset(
            {key for key, result in results.items() if "metrics" in result}
        )
        for results in (public_results, audit_results)
    )
    gate_passed = all(
        ranks[metric].get("korbyte-128k") == 1
        for ranks in (public_ranks, audit_ranks)
        for metric in ("fertility", "ebpb")
    ) and coverage_passed and all(
        metrics["exact_document_ratio"] == 1.0
        and metrics["unknown_token_ratio"] == 0.0
        for metrics in (ours_public, ours_audit)
    )
    report = {
        "schema_version": 1,
        "created_at": datetime.now(UTC).isoformat(),
        "definition": (
            "First place means rank 1 by both fertility and effective bits per byte "
            "among the successfully loaded, revision-pinned public comparison set."
        ),
        "datasets": {
            "public_korean": {
                "id": PUBLIC_BENCHMARK_ID,
                "revision": PUBLIC_BENCHMARK_REVISION,
                "domains": list(PUBLIC_DOMAINS),
                "selection_role": "development evaluation",
                "text_examples_redistributed": False,
                "results": public_results,
            },
            "kmmlu_test": {
                "id": KMMLU_ID,
                "revision": KMMLU_REVISION,
                "subjects": 45,
                "selection_role": "frozen post-selection audit; no algorithm changes followed",
                "text_examples_redistributed": False,
                "results": audit_results,
            },
        },
        "first_place_gate_passed": gate_passed,
        "comparison_coverage_gate_passed": coverage_passed,
        "required_successful_systems": sorted(required_systems),
        "limitations": [
            "Intrinsic tokenization metrics do not establish downstream model quality.",
            "The comparison set is broad and current but cannot prove a universal ranking.",
            "Thunder's public tokenizer uses an unsupported custom `Beta` model and "
            "is recorded as unavailable.",
        ],
        "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: importlib.metadata.version(package)
                for package in ("huggingface-hub", "rapidfuzz", "tokenizers", "transformers")
            },
        },
    }
    paths = Paths(root)
    paths.comparison_json.parent.mkdir(parents=True, exist_ok=True)
    paths.comparison_json.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return report