252 lines
10 KiB
Python
252 lines
10 KiB
Python
 
|
#!/usr/bin/env python3
|
||
|
|
"""Compare normalizer models on the same sentences.
|
||
|
|
|
||
|
|
The shipped model is a 0.6B task-specific fine-tune that needs no prompt but
|
||
|
|
miscounts large numbers and occasionally drops a word. This runs any causal LM
|
||
|
|
over the same battery so the trade is measurable rather than argued about:
|
||
|
|
|
||
|
|
python bench.py # the fine-tune
|
||
|
|
python bench.py --model models/PLLuM-4B-instruct-2512 --prompted
|
||
|
|
python bench.py --model … --prompted --save pllum.json
|
||
|
|
python bench.py --compare fine-tune.json pllum.json
|
||
|
|
|
||
|
|
`--prompted` switches from the fine-tune's bare "text in, text out" contract to
|
||
|
|
an instruction + few-shot prompt, which a general instruct model needs. Digit
|
||
|
|
tokens are banned from the output either way, and decoding is greedy, so a
|
||
|
|
difference between two runs is the model and not the sampling.
|
||
|
|
|
||
|
|
Automatic flags are cheap heuristics, not a grade: `digits` means an
|
||
|
|
un-normalized digit survived, `lost` means a word from the input has no
|
||
|
|
plausible inflected match in the output. Read the text before believing either.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
HERE = Path(__file__).resolve().parent
|
||
|
|
sys.path.insert(0, str(HERE))
|
||
|
|
|
||
|
|
import normalizer as nz # noqa: E402 (reuses its protect/restore and banned list)
|
||
|
|
|
||
|
|
SYSTEM_PROMPTED = (
|
||
|
|
"Jesteś narzędziem do normalizacji tekstu dla syntezatora mowy. "
|
||
|
|
"Zamieniasz wszystkie liczby zapisane cyframi na słowa, w poprawnej formie "
|
||
|
|
"gramatycznej wynikającej z kontekstu zdania. Nie zmieniasz, nie dodajesz "
|
||
|
|
"ani nie usuwasz żadnych innych słów. Zwracasz wyłącznie przekształcone "
|
||
|
|
"zdanie, bez komentarza."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Few-shot examples chosen to pin down the three things a general model gets
|
||
|
|
# wrong without them: the case must follow the governing word, the rest of the
|
||
|
|
# sentence must survive verbatim, and the answer must be the sentence alone.
|
||
|
|
FEWSHOT = [
|
||
|
|
("Koncert dla 8250 osób.",
|
||
|
|
"Koncert dla ośmiu tysięcy dwustu pięćdziesięciu osób."),
|
||
|
|
("Populacja Polski to 38 000 000 mieszkańców.",
|
||
|
|
"Populacja Polski to trzydzieści osiem milionów mieszkańców."),
|
||
|
|
("Mówimy o 25 psach.",
|
||
|
|
"Mówimy o dwudziestu pięciu psach."),
|
||
|
|
("Wygrał 1247 złotych.",
|
||
|
|
"Wygrał tysiąc dwieście czterdzieści siedem złotych."),
|
||
|
|
]
|
||
|
|
|
||
|
|
BATTERY = [
|
||
|
|
"Sprawa dotyczy 1 350 000 ludzi.",
|
||
|
|
"Koncert dla 8250 osób.",
|
||
|
|
"Budżet wynosi 12 500 000 złotych.",
|
||
|
|
"Populacja Polski to 38 000 000 mieszkańców.",
|
||
|
|
"Zapłacili 1 500 000 000 dolarów.",
|
||
|
|
"Zostało 999 999 sztuk.",
|
||
|
|
"Numer seryjny 1234567.",
|
||
|
|
"Ponad 100 000 widzów na stadionie.",
|
||
|
|
"Mówimy o 1 350 000 ludziach.",
|
||
|
|
"Z 8250 osobami nie da się dyskutować.",
|
||
|
|
"Firma zatrudnia 2 500 pracowników.",
|
||
|
|
"Dystans 42 195 metrów.",
|
||
|
|
"Cena spadła z 15 700 000 do 8 250 000.",
|
||
|
|
"Było tam 25 psów.",
|
||
|
|
"Mówimy o 25 psach.",
|
||
|
|
"Przyjechał z 25 psami.",
|
||
|
|
"Przyglądał się 25 psom.",
|
||
|
|
"Nie ma tu 25 psów.",
|
||
|
|
"Urodził się 3 maja 1791 roku.",
|
||
|
|
"Wzrost o 16 procent w 2026 roku.",
|
||
|
|
"Spotkanie o 14:30.",
|
||
|
|
"Godzina 23:56.",
|
||
|
|
"Radio 95,7 FM.",
|
||
|
|
"Około 3,5 metra.",
|
||
|
|
"Mieszka pod 12/45.",
|
||
|
|
"Wygrał 1247 złotych.",
|
||
|
|
"To kosztuje 1999 dolarów i 99 centów.",
|
||
|
|
"Mam 3 koty i 25 psów.",
|
||
|
|
"Zostało 5 minut.",
|
||
|
|
"Temperatura opon w porządku. Przód 85, tył 75.",
|
||
|
|
"Louis Chiron miał 57 lat, kiedy wygrał Grand Prix Monako.",
|
||
|
|
"Zwykłe zdanie bez żadnych liczb.",
|
||
|
|
]
|
||
|
|
|
||
|
|
_WORD_RE = re.compile(r'[^\W\d_]+', re.UNICODE)
|
||
|
|
|
||
|
|
|
||
|
|
def lost_words(src: str, out: str) -> list:
|
||
|
|
"""Input words with no plausible match in the output.
|
||
|
|
|
||
|
|
Polish inflects, so an exact comparison is useless — `ludziach` legitimately
|
||
|
|
becomes `ludziach` and nothing else, but `psach`/`psami` share a stem. A
|
||
|
|
five-character prefix is a crude stand-in for a lemma and is good enough to
|
||
|
|
notice a noun that vanished entirely, which is the failure worth catching.
|
||
|
|
"""
|
||
|
|
def stems(text):
|
||
|
|
return {w.lower()[:5] for w in _WORD_RE.findall(text)}
|
||
|
|
missing = stems(src) - stems(out)
|
||
|
|
return sorted(w for w in missing if len(w) >= 4)
|
||
|
|
|
||
|
|
|
||
|
|
def load_prompted(model_dir: str, gpu_gib: float = None,
|
||
|
|
system: str = None, fewshot: list = None):
|
||
|
|
"""A general instruct model: system prompt, few-shot, greedy, digits banned.
|
||
|
|
|
||
|
|
`gpu_gib` caps GPU use and spills the rest to RAM. Needed for the int4
|
||
|
|
"awq" builds: transformers decompresses compressed-tensors weights back to
|
||
|
|
bfloat16 at load, so an 11B model wants ~22 GB whatever the file size says.
|
||
|
|
Offloading is slow but it is the difference between measuring a model and
|
||
|
|
guessing about it.
|
||
|
|
"""
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
tok = AutoTokenizer.from_pretrained(model_dir)
|
||
|
|
kwargs = {"dtype": torch.bfloat16, "device_map": "auto"}
|
||
|
|
if gpu_gib == 0:
|
||
|
|
# Everything in RAM. Slow, but the only way to run a compressed-tensors
|
||
|
|
# int4 build in transformers: its decompression hook expands weights to
|
||
|
|
# bfloat16 in place, and it cannot read the meta tensors that accelerate
|
||
|
|
# leaves behind when max_memory forces a partial offload.
|
||
|
|
kwargs["device_map"] = {"": "cpu"}
|
||
|
|
elif gpu_gib:
|
||
|
|
kwargs["max_memory"] = {0: f"{gpu_gib}GiB", "cpu": "80GiB"}
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(model_dir, **kwargs).eval()
|
||
|
|
|
||
|
|
banned = set()
|
||
|
|
for tid in range(len(tok)):
|
||
|
|
piece = tok.decode([tid]).strip()
|
||
|
|
if piece and piece.isdigit():
|
||
|
|
banned.add(tid)
|
||
|
|
banned = sorted(banned)
|
||
|
|
print(f" {len(banned)} digit tokens banned; device "
|
||
|
|
f"{next(model.parameters()).device}", file=sys.stderr)
|
||
|
|
|
||
|
|
def run(text: str) -> str:
|
||
|
|
messages = [{"role": "system", "content": system or SYSTEM_PROMPTED}]
|
||
|
|
for user, assistant in (fewshot or FEWSHOT):
|
||
|
|
messages.append({"role": "user", "content": user})
|
||
|
|
messages.append({"role": "assistant", "content": assistant})
|
||
|
|
messages.append({"role": "user", "content": text})
|
||
|
|
prompt = tok.apply_chat_template(messages, tokenize=False,
|
||
|
|
add_generation_prompt=True)
|
||
|
|
inputs = tok(prompt, return_tensors="pt").to(model.device)
|
||
|
|
n = inputs["input_ids"].shape[1]
|
||
|
|
with torch.no_grad():
|
||
|
|
out = model.generate(**inputs, max_new_tokens=nz.MAX_NEW_TOKENS,
|
||
|
|
do_sample=False,
|
||
|
|
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||
|
|
suppress_tokens=banned)
|
||
|
|
answer = tok.decode(out[0][n:], skip_special_tokens=True).strip()
|
||
|
|
# Instruct models like to introduce themselves; keep the last non-empty
|
||
|
|
# line, which is the sentence in every observed case.
|
||
|
|
lines = [l.strip() for l in answer.splitlines() if l.strip()]
|
||
|
|
return lines[-1] if lines else answer
|
||
|
|
|
||
|
|
return run
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
ap = argparse.ArgumentParser(description=__doc__,
|
||
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
|
|
ap.add_argument("--model", default=None,
|
||
|
|
help="Model directory (default: the shipped fine-tune)")
|
||
|
|
ap.add_argument("--prompted", action="store_true",
|
||
|
|
help="Use an instruction + few-shot prompt (general models)")
|
||
|
|
ap.add_argument("--gpu-gib", type=float, default=None, metavar="N",
|
||
|
|
help="Cap GPU memory and spill the rest to RAM (slow, but "
|
||
|
|
"the only way to run a model that does not fit)")
|
||
|
|
ap.add_argument("--no-protect", action="store_true",
|
||
|
|
help="Do not shield standalone single digits — normalize the "
|
||
|
|
"text exactly as the model card does")
|
||
|
|
ap.add_argument("--save", metavar="FILE", help="Write results as JSON")
|
||
|
|
ap.add_argument("--compare", nargs=2, metavar=("A", "B"),
|
||
|
|
help="Print two saved runs side by side and exit")
|
||
|
|
args = ap.parse_args()
|
||
|
|
|
||
|
|
if args.compare:
|
||
|
|
a, b = (json.loads(Path(f).read_text(encoding="utf-8")) for f in args.compare)
|
||
|
|
agree = 0
|
||
|
|
for ra, rb in zip(a["results"], b["results"]):
|
||
|
|
if ra["out"] == rb["out"]:
|
||
|
|
agree += 1
|
||
|
|
continue
|
||
|
|
print(f" in : {ra['in']}")
|
||
|
|
print(f" A : {ra['out']}")
|
||
|
|
print(f" B : {rb['out']}\n")
|
||
|
|
print(f"identical: {agree}/{len(a['results'])} "
|
||
|
|
f"A={a['model']} B={b['model']}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
if args.prompted:
|
||
|
|
if not args.model:
|
||
|
|
ap.error("--prompted needs --model")
|
||
|
|
run = load_prompted(args.model, args.gpu_gib)
|
||
|
|
label = args.model
|
||
|
|
else:
|
||
|
|
if args.model:
|
||
|
|
nz.MODEL_DIR = Path(args.model)
|
||
|
|
if not nz.load_model():
|
||
|
|
print("model unavailable; refusing to benchmark the fallback",
|
||
|
|
file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
run = nz.normalize
|
||
|
|
label = str(nz.MODEL_DIR)
|
||
|
|
if args.no_protect:
|
||
|
|
nz.KEEP_RE = re.compile(r'(?!x)x') # matches nothing
|
||
|
|
label += " [no single-digit protection]"
|
||
|
|
|
||
|
|
results, t0 = [], time.time()
|
||
|
|
for text in BATTERY:
|
||
|
|
t = time.time()
|
||
|
|
out = run(text)
|
||
|
|
results.append({"in": text, "out": out, "sec": round(time.time() - t, 2),
|
||
|
|
"digits": bool(nz.DIGIT_RE.search(nz.KEEP_RE.sub("", out))),
|
||
|
|
"lost": lost_words(text, out)})
|
||
|
|
total = time.time() - t0
|
||
|
|
|
||
|
|
for r in results:
|
||
|
|
flags = []
|
||
|
|
if r["digits"]:
|
||
|
|
flags.append("DIGITS LEFT")
|
||
|
|
if r["lost"]:
|
||
|
|
flags.append("LOST: " + ", ".join(r["lost"]))
|
||
|
|
print(f" in : {r['in']}")
|
||
|
|
print(f" out: {r['out']}" + (f" [{'; '.join(flags)}]" if flags else ""))
|
||
|
|
print()
|
||
|
|
|
||
|
|
n_dig = sum(r["digits"] for r in results)
|
||
|
|
n_lost = sum(bool(r["lost"]) for r in results)
|
||
|
|
print(f"{label}\n {len(results)} sentences in {total:.1f}s "
|
||
|
|
f"({total/len(results):.2f}s each)\n"
|
||
|
|
f" {n_dig} with digits left, {n_lost} with a word apparently lost")
|
||
|
|
|
||
|
|
if args.save:
|
||
|
|
Path(args.save).write_text(json.dumps(
|
||
|
|
{"model": label, "prompted": args.prompted, "seconds": round(total, 1),
|
||
|
|
"results": results}, ensure_ascii=False, indent=1), encoding="utf-8")
|
||
|
|
print(f" saved to {args.save}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|