Files
model-training/llm/decision_gate.py
T
2026-07-19 23:52:25 +04:00

62 lines
2.3 KiB
Python

"""Compare raw and CPT evaluation JSON and enforce the post-CPT gate."""
import argparse
import json
import re
import sys
from pathlib import Path
def repetition_4gram_rate(report: dict) -> float:
rates = []
for lang in ("ru", "en"):
for item in report["generation"].get(lang, []):
words = re.findall(r"\w+", item["output"].lower())
grams = [tuple(words[i:i + 4]) for i in range(max(0, len(words) - 3))]
rates.append(1 - len(set(grams)) / len(grams) if grams else 0.0)
return sum(rates) / len(rates) if rates else 1.0
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--raw", required=True)
ap.add_argument("--cpt", required=True)
ap.add_argument("--max-ru-ppl-ratio", type=float, default=0.98)
ap.add_argument("--max-en-ppl-ratio", type=float, default=1.10)
ap.add_argument("--max-cyrillic-drop", type=float, default=0.0)
args = ap.parse_args()
raw = json.loads(Path(args.raw).read_text())
cpt = json.loads(Path(args.cpt).read_text())
raw_repetition = repetition_4gram_rate(raw)
cpt_repetition = repetition_4gram_rate(cpt)
checks = {
"ru_ppl_improved_by_2pct": cpt["perplexity"]["ru"] <= raw["perplexity"]["ru"] * args.max_ru_ppl_ratio,
"en_ppl_regression_within_10pct": cpt["perplexity"]["en"] <= raw["perplexity"]["en"] * args.max_en_ppl_ratio,
"cyrillic_validity_not_worse": cpt["generation"]["ru_valid_pct"] >= raw["generation"]["ru_valid_pct"] - args.max_cyrillic_drop,
"english_retained": bool(cpt["generation"]["en_retained"]),
"generation_repetition_not_degraded": (
cpt_repetition <= max(0.15, raw_repetition + 0.05)
),
}
report = {
"pass": all(checks.values()),
"checks": checks,
"raw": raw.get("model"),
"cpt": cpt.get("model"),
"ratios": {
"ru_ppl": cpt["perplexity"]["ru"] / raw["perplexity"]["ru"],
"en_ppl": cpt["perplexity"]["en"] / raw["perplexity"]["en"],
},
"generation_repetition_4gram_rate": {
"raw": raw_repetition,
"cpt": cpt_repetition,
"maximum_allowed": max(0.15, raw_repetition + 0.05),
},
}
print(json.dumps(report, ensure_ascii=False, indent=2))
sys.exit(0 if report["pass"] else 1)
if __name__ == "__main__":
main()