1# eval_sage.py [out_file]
 2# Quantitative quality check of SAGE full-text correction on a battery of realistic Russian
 3# typo sentences. Lenient match (normalize spaces/ё/case/trailing punctuation, substring ok),
 4# since SAGE also adds punctuation & capitalization. Writes results (UTF-8) to out_file.
 5import sys, re
 6import torch
 7from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 8
 9MODEL = "ai-forever/sage-fredt5-distilled-95m"
10tok = AutoTokenizer.from_pretrained(MODEL)
11model = AutoModelForSeq2SeqLM.from_pretrained(MODEL)
12model.eval()
13
14def correct(t):
15    ids = tok(t, return_tensors="pt", max_length=256, truncation=True)
16    with torch.no_grad():
17        out = model.generate(**ids, max_length=256, num_beams=5, early_stopping=True)
18    return tok.batch_decode(out, skip_special_tokens=True)[0]
19
20def norm(s):
21    s = re.sub(r"\s+", " ", s.strip().lower().replace("ё", "е"))
22    return s.strip(" .!?,")
23
24TESTS = [
25    ("сегодя я хочю праверить арфографию маей пограммы", "сегодня я хочу проверить орфографию моей программы"),
26    ("он сказал што кампьютер не работаит и нужно перезагрузить", "он сказал что компьютер не работает и нужно перезагрузить"),
27    ("мы пашли в магазин и купили малако хлеб и агурцы", "мы пошли в магазин и купили молоко хлеб и огурцы"),
28    ("здраствуйте дарогие друзя как ваши дила", "здравствуйте дорогие друзья как ваши дела"),
29    ("я незнаю што мне делать в этай ситуацыи", "я не знаю что мне делать в этой ситуации"),
30    ("он абизательно прийдет завтра утрам", "он обязательно придет завтра утром"),
31    ("это очень харошая и интиресная идея", "это очень хорошая и интересная идея"),
32    ("мы должны зделать это сегодня вечиром", "мы должны сделать это сегодня вечером"),
33    ("впринципе мне нравитца этот падход", "в принципе мне нравится этот подход"),
34    ("навернае ты прав наш чёт", "наверное ты прав"),
35    ("чтобы добится успеха нужно много работать", "чтобы добиться успеха нужно много работать"),
36    ("извените за апоздание я был занят", "извините за опоздание я был занят"),
37    ("сколько стоет билет на поизд до москвы", "сколько стоит билет на поезд до москвы"),
38    ("painаю что нужно но немагу", "понимаю что нужно но не могу"),
39    ("ришил написать тибе письмо", "решил написать тебе письмо"),
40    ("прачитай этот тextст внимательна", "прочитай этот текст внимательно"),
41]
42
43out_path = sys.argv[1] if len(sys.argv) > 1 else "Spell/_eval_out.txt"
44ok = 0
45lines = []
46for inp, exp in TESTS:
47    r = correct(inp)
48    passed = (norm(exp) == norm(r)) or (norm(exp) in norm(r))
49    if passed:
50        ok += 1
51    lines.append(("PASS " if passed else "FAIL ") + repr(inp) + " -> " + repr(r) +
52                 ("" if passed else "  (want ~ " + repr(exp) + ")"))
53lines.append("")
54lines.append("=== %d/%d = %d%% (lenient match) ===" % (ok, len(TESTS), 100 * ok // len(TESTS)))
55with open(out_path, "w", encoding="utf-8") as f:
56    f.write("\n".join(lines))