code/eos_period.py

#!/usr/bin/env python3
"""LE POINT FINAL EN TROP — 2e ecart revele par l analyse des CSV.

                        reference train   notre sortie
  finit par "." lin          67.2 %          96.0 %
  finit par "." sna          95.1 %          99.8 %

=> ~128 clips lingala portent un point final que la reference n a pas. Chacun coute
1 substitution de mot (dernier mot "mot." au lieu de "mot") ET 1 insertion de caractere.
Le scorer est BRUT : la ponctuation compte.

Meme cause que les virgules : le CTC a appris a terminer par un point parce que c est
le cas le plus frequent, sans signal acoustique pour trancher. Le shona est presque
juste (95.1 -> 99.8), le LINGALA est le vrai gisement (67.2 -> 96.0).

Ce script fait TOUT : entraine un classifieur de phrase "cet enonce se termine-t-il
par un point ?" sur les transcriptions, puis RETIRE le point final des clips ou le
modele en est sûr. Augmentation par bruit ASR conservee (elle avait fait passer la
precision des virgules de 0.48 a 0.70).

SEUIL = probabilite MINIMALE de "pas de point" requise pour retirer le point.
SEUIL > 1 => aucune suppression => reproduit BASE a l identique (CONTROLE).
On ne fait que RETIRER : on n ajoute jamais de point.
"""
import argparse
import csv
import json
import os
import random

import numpy as np
import torch
from datasets import Dataset
from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
                          DataCollatorWithPadding, Trainer, TrainingArguments)


def parse():
    p = argparse.ArgumentParser()
    p.add_argument("--base", default="Davlan/afro-xlmr-base")
    p.add_argument("--out", default="/scratch/runs/eos")
    p.add_argument("--train", nargs="+", required=True)
    p.add_argument("--sub", default="/root/sub_XC010.csv")
    p.add_argument("--langf", default="/root/test_lang.json")
    p.add_argument("--ths", default="1.01,0.9,0.8,0.7,0.6")
    p.add_argument("--tag", default="EO")
    p.add_argument("--lr", type=float, default=3e-5)
    p.add_argument("--epochs", type=float, default=3)
    p.add_argument("--bs", type=int, default=32)
    p.add_argument("--maxlen", type=int, default=192)
    p.add_argument("--noise_copies", type=int, default=2)
    p.add_argument("--noise_rate", type=float, default=0.35)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def corrupt(text, rng, rate):
    out = []
    for w in text.split():
        if rng.random() < rate and len(w) > 2:
            i = rng.randrange(len(w))
            r = rng.random()
            if r < 0.45:
                w = w[:i] + rng.choice("aeioubkmnlstz") + w[i + 1:]
            elif r < 0.8:
                w = w[:i] + w[i + 1:]
            else:
                w = w[:i] + rng.choice("aeiounm") + w[i:]
        out.append(w or "a")
    return " ".join(out)


def strip_final(t):
    """retire le point final eventuel ; renvoie (texte_sans, avait_un_point)"""
    t = t.rstrip()
    if t.endswith("."):
        return t[:-1].rstrip(), 1
    return t, 0


def main():
    a = parse()
    os.makedirs(a.out, exist_ok=True)
    random.seed(a.seed)
    torch.manual_seed(a.seed)

    rows = []
    for m in a.train:
        for line in open(m, encoding="utf-8"):
            t = " ".join(str(json.loads(line).get("text", "")).split())
            if len(t.split()) < 3:
                continue
            body, has = strip_final(t)
            if body:
                rows.append({"text": body, "label": has})
    rng = random.Random(a.seed + 7)
    if a.noise_copies:
        aug = [{"text": corrupt(r["text"], rng, a.noise_rate * (c + 1) / a.noise_copies),
                "label": r["label"]}
               for r in rows for c in range(a.noise_copies)]
        rows = rows + aug
    random.shuffle(rows)
    nv = max(500, len(rows) // 20)
    ev, tr = rows[:nv], rows[nv:]
    print("TRAIN %d | EVAL %d | taux de point final %.1f%%"
          % (len(tr), len(ev), 100 * np.mean([r["label"] for r in tr])), flush=True)

    tok = AutoTokenizer.from_pretrained(a.base)
    model = AutoModelForSequenceClassification.from_pretrained(a.base, num_labels=2)

    def enc(b):
        return tok(b["text"], truncation=True, max_length=a.maxlen)

    dtr = Dataset.from_list(tr).map(enc, batched=True, remove_columns=["text"])
    dev = Dataset.from_list(ev).map(enc, batched=True, remove_columns=["text"])

    def metrics(p):
        pred = np.argmax(p[0], -1)
        la = p[1]
        # classe utile = 0 ("pas de point"), c est elle qu on va agir
        tp = int(((pred == 0) & (la == 0)).sum())
        fp = int(((pred == 0) & (la == 1)).sum())
        fn = int(((pred == 1) & (la == 0)).sum())
        pr = tp / max(tp + fp, 1)
        rc = tp / max(tp + fn, 1)
        return {"acc": float((pred == la).mean()), "precision_sans_point": pr,
                "rappel_sans_point": rc,
                "f1": 2 * pr * rc / max(pr + rc, 1e-9)}

    args = TrainingArguments(
        output_dir=a.out, per_device_train_batch_size=a.bs,
        per_device_eval_batch_size=128, num_train_epochs=a.epochs,
        learning_rate=a.lr, warmup_ratio=0.1, bf16=True,
        eval_strategy="epoch", save_strategy="epoch", save_total_limit=2,
        load_best_model_at_end=True, metric_for_best_model="f1",
        greater_is_better=True, logging_steps=200, report_to=[],
        dataloader_num_workers=4, seed=a.seed)
    trainer = Trainer(model=model, args=args, train_dataset=dtr, eval_dataset=dev,
                      data_collator=DataCollatorWithPadding(tok), compute_metrics=metrics)
    trainer.train()
    print("EVAL:", json.dumps(trainer.evaluate(), indent=1), flush=True)

    # ---- courbe precision/seuil sur les donnees retenues : DECISIF avant de soumettre
    pr_ev = trainer.predict(dev).predictions
    P0 = torch.softmax(torch.tensor(pr_ev).float(), -1)[:, 0].numpy()
    Y = np.array([r["label"] for r in ev])
    print("seuil | suppressions | precision | rappel | gain net")
    for th in (0.9, 0.8, 0.7, 0.6, 0.5):
        s = P0 >= th
        tp = int((s & (Y == 0)).sum())
        fp = int((s & (Y == 1)).sum())
        print("%.2f  | %7d      |   %.3f   | %.3f  | %+d %s"
              % (th, int(s.sum()), tp / max(tp + fp, 1), tp / max(int((Y == 0).sum()), 1),
                 tp - fp, "OK" if tp > fp else "PERDANT"), flush=True)

    # ---- application a la soumission
    base = {r["ID"]: r["Target"] for r in csv.DictReader(open(a.sub, encoding="utf-8"))}
    lang = json.load(open(a.langf))
    ids = list(base)
    bodies, hasdot = {}, {}
    for k in ids:
        b, h = strip_final(str(base[k]))
        bodies[k], hasdot[k] = b, h
    model.eval()
    PROB = {}
    with torch.inference_mode():
        for i in range(0, len(ids), 64):
            ch = ids[i:i + 64]
            x = tok([bodies[k] or "a" for k in ch], truncation=True, max_length=a.maxlen,
                    padding=True, return_tensors="pt")
            lo = model(**{kk: vv.cuda() for kk, vv in x.items()}).logits.float()
            p0 = torch.softmax(lo, -1)[:, 0].cpu().numpy()
            for j, k in enumerate(ch):
                PROB[k] = float(p0[j])

    from huggingface_hub import HfApi
    api = HfApi(token=open(os.path.expanduser("~/.cache/huggingface/token")).read().strip())
    for th in [float(x) for x in a.ths.split(",")]:
        out = dict(base)
        nrem = {"lin": 0, "sna": 0}
        ncl = {"lin": 0, "sna": 0}
        for k in ids:
            lg = lang.get(k, "lin")
            ncl[lg] = ncl.get(lg, 0) + 1
            if hasdot[k] and PROB[k] >= th and bodies[k]:
                out[k] = bodies[k]
                nrem[lg] = nrem.get(lg, 0) + 1
        nchg = sum(1 for k in ids if out[k] != base[k])
        empt = sum(1 for x in out.values() if not str(x).strip())
        tag = "%s%03d" % (a.tag, round(th * 100))
        OUT = "/root/sub_%s.csv" % tag
        with open(OUT, "w", newline="", encoding="utf-8") as f:
            wr = csv.writer(f)
            wr.writerow(["ID", "Target"])
            for k in base:
                wr.writerow([k, out[k] or "a"])
        assert len(out) == 892 and empt == 0, "%s INVALIDE" % tag
        pl = 1 - nrem["lin"] / max(ncl["lin"], 1) * 1.0
        rl = sum(1 for k in ids if lang.get(k) == "lin" and out[k].rstrip().endswith(".")) \
            / max(ncl["lin"], 1)
        rs = sum(1 for k in ids if lang.get(k) == "sna" and out[k].rstrip().endswith(".")) \
            / max(ncl["sna"], 1)
        if th <= 1.0:
            api.upload_file(path_or_fileobj=OUT,
                            path_in_repo="phase2_corrected/sub_%s.csv" % tag,
                            repo_id="Pricile/waxal2026-backup", repo_type="model")
        flag = "  <-- CONTROLE : doit etre 0" if th > 1.0 else ""
        print("%-7s seuil %.2f | clips modifies %3d/892 | fin. lin %.1f%% (ref 67.2) "
              "sna %.1f%% (ref 95.1)%s" % (tag, th, nchg, 100 * rl, 100 * rs, flag), flush=True)
    print("EOS_DONE", flush=True)


if __name__ == "__main__":
    main()