| |
| """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() |
|
|
