9bb00d49f2
TextMachine project repository (AI translation of literary books). Includes: v2 architecture decisions, MVP plan, research 01-12, polygon experiments 01-03, backend-session revalidation verdict (03-implementation-notes.md). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
206 lines
8.7 KiB
Python
206 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
||
"""Токен-калькулятор TextMachine: замер фактической токенизации художественных текстов.
|
||
|
||
Считает для каждого файла eval/data/samples/<lang>/*.txt метрики по всем токенизаторам:
|
||
- tokens/CJK-знак (zh, ja), tokens/слово (en, ru), символов/токен;
|
||
- при наличии manifest.json с parallel_of — соотношение объёмов перевод/оригинал.
|
||
|
||
Запуск: eval/.venv/bin/python eval/token_calc.py [--samples-dir eval/data/samples] [--json-out ...]
|
||
Выход: markdown-таблицы в stdout + JSON с сырыми числами.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import sys
|
||
import unicodedata
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
|
||
# --- Токенизаторы -----------------------------------------------------------
|
||
|
||
def load_tokenizers() -> dict[str, "callable"]:
|
||
"""Возвращает {имя: fn(text)->int}. Каждый fn считает токены."""
|
||
toks: dict[str, callable] = {}
|
||
|
||
import tiktoken
|
||
for enc_name, label in [("o200k_base", "gpt-o200k (GPT-4o/5)"),
|
||
("cl100k_base", "gpt-cl100k (GPT-4)")]:
|
||
enc = tiktoken.get_encoding(enc_name)
|
||
toks[label] = (lambda e: lambda t: len(e.encode(t, disallowed_special=())))(enc)
|
||
|
||
from tokenizers import Tokenizer
|
||
hf_dirs = {
|
||
"deepseek-v3.2": "DeepSeek V3.2",
|
||
"qwen3": "Qwen3",
|
||
"qwen2.5": "Qwen2.5",
|
||
"glm-4.6": "GLM-4.6",
|
||
}
|
||
for d, label in hf_dirs.items():
|
||
path = ROOT / "tokenizers" / d / "tokenizer.json"
|
||
if path.exists():
|
||
tk = Tokenizer.from_file(str(path))
|
||
toks[label] = (lambda t_: lambda t: len(t_.encode(t, add_special_tokens=False).ids))(tk)
|
||
else:
|
||
print(f"[warn] нет {path}, пропускаю {label}", file=sys.stderr)
|
||
return toks
|
||
|
||
|
||
# --- Подсчёт символов/слов ---------------------------------------------------
|
||
|
||
CJK_RE = re.compile(
|
||
"[㐀-䶿一-鿿豈-\U00020000-\U0002ebef]"
|
||
) # хань-иероглифы (кандзи/ханьцзы), вкл. расширения
|
||
KANA_RE = re.compile("[-ヿ]") # хирагана + катакана
|
||
WORD_RE = re.compile(r"[A-Za-zА-Яа-яЁёÀ-ɏ]+(?:[-'’][A-Za-zА-Яа-яЁё]+)*")
|
||
|
||
|
||
@dataclass
|
||
class TextStats:
|
||
path: str
|
||
lang: str
|
||
chars_total: int = 0
|
||
chars_no_ws: int = 0
|
||
cjk_chars: int = 0 # хань-иероглифы
|
||
kana_chars: int = 0 # ja: кана
|
||
words: int = 0 # en/ru: слова
|
||
tokens: dict = field(default_factory=dict) # {tokenizer: n}
|
||
|
||
|
||
def analyze(path: Path, lang: str, toks: dict) -> TextStats:
|
||
text = path.read_text(encoding="utf-8")
|
||
# нормализуем переводы строк, убираем BOM
|
||
text = text.replace("", "").strip()
|
||
st = TextStats(path=str(path), lang=lang)
|
||
st.chars_total = len(text)
|
||
st.chars_no_ws = sum(1 for c in text if not c.isspace())
|
||
st.cjk_chars = len(CJK_RE.findall(text))
|
||
st.kana_chars = len(KANA_RE.findall(text))
|
||
st.words = len(WORD_RE.findall(text))
|
||
for name, fn in toks.items():
|
||
st.tokens[name] = fn(text)
|
||
return st
|
||
|
||
|
||
def unit_count(st: TextStats) -> tuple[str, int]:
|
||
"""Опорная единица языка: zh/ja — CJK-знаки (для ja — кандзи+кана), en/ru — слова."""
|
||
if st.lang == "zh":
|
||
return "иероглиф", st.cjk_chars
|
||
if st.lang == "ja":
|
||
return "знак (кандзи+кана)", st.cjk_chars + st.kana_chars
|
||
return "слово", st.words
|
||
|
||
|
||
# --- Отчёт -------------------------------------------------------------------
|
||
|
||
def fmt_table(rows: list[list[str]], header: list[str]) -> str:
|
||
out = ["| " + " | ".join(header) + " |",
|
||
"|" + "|".join("---" for _ in header) + "|"]
|
||
out += ["| " + " | ".join(r) + " |" for r in rows]
|
||
return "\n".join(out)
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--samples-dir", default=str(ROOT / "data" / "samples"))
|
||
ap.add_argument("--json-out", default=str(ROOT / "data" / "token_calc_results.json"))
|
||
args = ap.parse_args()
|
||
|
||
samples = Path(args.samples_dir)
|
||
toks = load_tokenizers()
|
||
tok_names = list(toks)
|
||
|
||
stats: list[TextStats] = []
|
||
for lang_dir in sorted(samples.iterdir()):
|
||
if not lang_dir.is_dir():
|
||
continue
|
||
lang = lang_dir.name
|
||
for f in sorted(lang_dir.glob("*.txt")):
|
||
stats.append(analyze(f, lang, toks))
|
||
|
||
if not stats:
|
||
sys.exit(f"Нет файлов в {samples}/<lang>/*.txt")
|
||
|
||
# 1. Таблица per-file: tokens/единица по каждому токенизатору
|
||
print("## Токены на опорную единицу языка (zh/ja — знак, en/ru — слово)\n")
|
||
header = ["Файл", "Язык", "Единиц"] + tok_names
|
||
rows = []
|
||
for st in stats:
|
||
unit, n = unit_count(st)
|
||
rows.append([Path(st.path).name, st.lang, f"{n}"]
|
||
+ [f"{st.tokens[t] / n:.3f}" if n else "—" for t in tok_names])
|
||
print(fmt_table(rows, header))
|
||
|
||
# 2. Агрегат по языку (взвешенно по объёму)
|
||
print("\n## Агрегат по языку (сумма токенов / сумма единиц)\n")
|
||
rows = []
|
||
langs = sorted({st.lang for st in stats})
|
||
agg: dict[str, dict] = {}
|
||
for lang in langs:
|
||
group = [st for st in stats if st.lang == lang]
|
||
n_units = sum(unit_count(st)[1] for st in group)
|
||
n_chars = sum(st.chars_no_ws for st in group)
|
||
row = [lang, unit_count(group[0])[0], f"{n_units}"]
|
||
agg[lang] = {"unit": unit_count(group[0])[0], "units": n_units,
|
||
"chars_no_ws": n_chars, "per_unit": {}, "chars_per_token": {}}
|
||
for t in tok_names:
|
||
total = sum(st.tokens[t] for st in group)
|
||
row.append(f"{total / n_units:.3f}")
|
||
agg[lang]["per_unit"][t] = total / n_units
|
||
agg[lang]["chars_per_token"][t] = n_chars / total
|
||
rows.append(row)
|
||
print(fmt_table(rows, ["Язык", "Единица", "Единиц"] + tok_names))
|
||
|
||
print("\n## Символов (без пробелов) на токен\n")
|
||
rows = [[lang, *(f"{agg[lang]['chars_per_token'][t]:.2f}" for t in tok_names)]
|
||
for lang in langs]
|
||
print(fmt_table(rows, ["Язык"] + tok_names))
|
||
|
||
# 3. Параллельные пары из manifest.json
|
||
manifest_path = samples / "manifest.json"
|
||
pairs_out = []
|
||
if manifest_path.exists():
|
||
manifest = json.loads(manifest_path.read_text())
|
||
by_path = {st.path: st for st in stats}
|
||
entries = manifest if isinstance(manifest, list) else manifest.get("files", [])
|
||
pairs = [(e["path"], e["parallel_of"]) for e in entries if e.get("parallel_of")]
|
||
if pairs:
|
||
print("\n## Параллельные пары: объём перевода к оригиналу (в токенах одного токенизатора)\n")
|
||
rows = []
|
||
|
||
def resolve(p):
|
||
return by_path.get(str(ROOT.parent / p)) or by_path.get(p)
|
||
|
||
for tr_path, src_spec in pairs:
|
||
# parallel_of может быть списком файлов (перевод покрывает несколько)
|
||
src_paths = src_spec if isinstance(src_spec, list) else [src_spec]
|
||
tr = resolve(tr_path)
|
||
srcs = [resolve(s) for s in src_paths]
|
||
if not tr or not all(srcs):
|
||
print(f"[warn] пара не найдена: {tr_path} / {src_spec}", file=sys.stderr)
|
||
continue
|
||
src_tokens = {t: sum(s.tokens[t] for s in srcs) for t in tok_names}
|
||
ratios = {t: tr.tokens[t] / src_tokens[t] for t in tok_names}
|
||
pairs_out.append({"translation": tr.path, "source": src_paths,
|
||
"src_lang": srcs[0].lang, "ratios": ratios})
|
||
rows.append([Path(tr.path).name,
|
||
"+".join(Path(s).name for s in src_paths)]
|
||
+ [f"{ratios[t]:.2f}" for t in tok_names])
|
||
if rows:
|
||
print(fmt_table(rows, ["Перевод", "Оригинал"] + tok_names))
|
||
|
||
out = {
|
||
"files": [vars(st) for st in stats],
|
||
"aggregate_by_lang": agg,
|
||
"parallel_pairs": pairs_out,
|
||
"tokenizers": tok_names,
|
||
}
|
||
Path(args.json_out).write_text(json.dumps(out, ensure_ascii=False, indent=1))
|
||
print(f"\nJSON: {args.json_out}", file=sys.stderr)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|