rababa/eval_urdu_comparable.py at main · interscript/rababa · GitHub

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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
"""Urdu comparable eval — settle urd-diac-1.0 vs d2 on ONE harness.

urd-diac-1.0 (urdu_diacrit/run-001, ByT5-small) shipped with 3.74 CER
on urdu-diacrit/test.jsonl; d2 (rababa_urdu_byt5/run-002-d2,
ByT5-base) measured 5.77 within its own lineage. Different lineages,
possibly different test files — this runs BOTH models greedy on the
SAME urdu-diacrit/test.jsonl so the manifest can name the real best.

Usage:
    modal run --detach eval_urdu_comparable.py
"""

from __future__ import annotations

import modal

urdu_volume = modal.Volume.from_name("urdu-diacrit-datasets", create_if_missing=False)
udiac_volume = modal.Volume.from_name("urdu-diacrit-checkpoints", create_if_missing=False)
checkpoints_volume = modal.Volume.from_name("rababa-checkpoints", create_if_missing=True)

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("torch==2.5.1", "transformers==4.46.3", "editdistance", "tqdm")
)

# the shipped urd-diac-1.0 checkpoint was saved/validated under the
# export stack (torch 2.12.1 / transformers 5.14.1); under 4.46.3 it
# generates EMPTY strings. evaluate each model under its own stack.
image514 = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("torch==2.12.1", "transformers==5.14.1", "editdistance", "tqdm")
)

app = modal.App("rababa-urdu-comparable", image=image)

MODELS = {
    "urd_diacit_run001_shipped": ("/volumes/udiac/urdu_diacrit/run-001/best", image514),
    "d2_rababa_urdu_byt5": ("/volumes/ckpts/rababa_urdu_byt5/run-002-d2/best", image),
}


@app.function(
    gpu="A10G",
    timeout=2 * 60 * 60,
    volumes={
        "/datasets": urdu_volume,
        "/volumes/udiac": udiac_volume,
        "/volumes/ckpts": checkpoints_volume,
    },
)
def load_pairs() -> list[list[str]]:
    import json

    pairs = []
    with open("/datasets/urdu-diacrit/test.jsonl", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                row = json.loads(line)
                if row.get("src") and row.get("tgt"):
                    pairs.append([row["src"].strip(), row["tgt"].strip()])
    print(f"[data] {len(pairs)} test pairs", flush=True)
    return pairs


@app.function(
    gpu="A10G",
    timeout=2 * 60 * 60,
    image=image514,
    volumes={
        "/datasets": urdu_volume,
        "/volumes/udiac": udiac_volume,
        "/volumes/ckpts": checkpoints_volume,
    },
)
def eval_model(ckpt: str) -> dict:
    import editdistance
    import torch
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

    for v in (urdu_volume, udiac_volume, checkpoints_volume):
        v.reload()

    pairs = [(s, t) for s, t in load_pairs.remote()]
    tok = AutoTokenizer.from_pretrained(ckpt)
    model = AutoModelForSeq2SeqLM.from_pretrained(ckpt).to("cuda")
    model.eval()
    preds: list[str] = []
    with torch.no_grad():
        for i in range(0, len(pairs), 64):
            chunk = pairs[i : i + 64]
            enc = tok([s for s, _ in chunk], return_tensors="pt", padding=True,
                      truncation=True, max_length=256).to("cuda")
            gen = model.generate(**enc, max_new_tokens=256, num_beams=1)
            preds.extend(tok.batch_decode(gen, skip_special_tokens=True))
    total_ed = total_len = exact = 0
    for (_, tgt), pred in zip(pairs, preds):
        total_ed += editdistance.eval(pred.strip(), tgt)
        total_len += len(tgt)
        exact += int(pred.strip() == tgt)
    res = {"cer": 100 * total_ed / max(1, total_len), "word_acc": 100 * exact / len(pairs)}
    print(f"[{ckpt}] CER={res['cer']:.2f} word_acc={res['word_acc']:.2f}", flush=True)
    return res


@app.function(
    gpu="A10G",
    timeout=2 * 60 * 60,
    volumes={
        "/datasets": urdu_volume,
        "/volumes/udiac": udiac_volume,
        "/volumes/ckpts": checkpoints_volume,
    },
)
def evaluate() -> dict:
    return {name: eval_model.remote(ckpt) for name, (ckpt, _img) in MODELS.items()}


@app.local_entrypoint()
def main():
    import json

    print(json.dumps(evaluate.remote(), indent=2))