data_studio/tokenizers.py

from __future__ import annotations

from typing import Any


REQUIRED_TOKENIZER_METRICS = frozenset({"telugu_fertility", "compression_ratio", "unknown_rate", "downstream_loss"})


def tokenizer_candidate_report(
    tokenizer_version: str,
    vocabulary_size: int,
    algorithm: str,
    metrics: dict[str, float],
) -> dict[str, Any]:
    missing = sorted(REQUIRED_TOKENIZER_METRICS.difference(metrics))
    if missing:
        raise ValueError(f"Tokenizer report is missing metrics: {missing}")
    return {
        "schema_version": "1.0",
        "tokenizer_version": tokenizer_version,
        "vocabulary_size": vocabulary_size,
        "algorithm": algorithm,
        "metrics": metrics,
        "status": "candidate",
    }


def compare_tokenizer_candidates(candidates: list[dict[str, Any]]) -> dict[str, Any]:
    if len(candidates) < 2:
        raise ValueError("Tokenizer promotion requires at least two candidates.")
    versions = [candidate["tokenizer_version"] for candidate in candidates]
    return {"candidate_versions": versions, "promotion_requires_owner_decision": True}