File size: 2,136 Bytes
6eed659 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | #!/usr/bin/env python3
"""Applique le restaurateur de ponctuation a un CSV de soumission (ID,Target).
Garde anti-regression : rejette toute sortie qui modifie les mots hors ponct/casse.
Langue deduite du prefixe de l'ID (lin_/lug_/sna_)."""
import argparse
import csv
import unicodedata
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration
MODEL = "/root/models/punct_best"
def strip_form(t):
t = unicodedata.normalize("NFC", t).lower()
t = "".join(c if c.isalnum() or c.isspace() or c in "'-" else " " for c in t)
return " ".join(t.split())
@torch.inference_mode()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--inp", required=True)
ap.add_argument("--out", required=True)
a = ap.parse_args()
tok = AutoTokenizer.from_pretrained(MODEL)
model = T5ForConditionalGeneration.from_pretrained(MODEL, torch_dtype=torch.bfloat16).cuda().eval()
rows = list(csv.DictReader(open(a.inp, encoding="utf-8")))
kept = rej = 0
B = 48
out_map = {}
for k in range(0, len(rows), B):
chunk = rows[k:k+B]
srcs, langs = [], []
for r in chunk:
lang = r["ID"].split("_")[0]
langs.append(lang)
srcs.append(f"{lang}: {strip_form(r['Target'])}")
enc = tok(srcs, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda")
gen = model.generate(**enc, max_length=512, num_beams=1)
outs = tok.batch_decode(gen, skip_special_tokens=True)
for r, o in zip(chunk, outs):
o = " ".join(o.split())
if o and strip_form(o) == strip_form(r["Target"]):
out_map[r["ID"]] = o
kept += 1
else:
out_map[r["ID"]] = r["Target"] or "a"
rej += 1
with open(a.out, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow(["ID", "Target"])
for r in rows:
w.writerow([r["ID"], out_map[r["ID"]] or "a"])
print(f"PUNCT_CSV_DONE {a.out}: {kept} restaures, {rej} gardes-original")
if __name__ == "__main__":
main()
|