3333 "transformers==5.14.1" ,
3434 "pyyaml>=6.0" ,
3535 "numpy>=1.26" ,
36+ # arabic gate harness: Misraj evaluator + SadeedDiac-25 parquet
37+ "pyarabic" ,
38+ "pandas" ,
39+ "pyarrow" ,
3640 )
3741 .env ({"PYTORCH_CUDA_ALLOC_CONF" : "expandable_segments:True" })
3842 .add_local_dir (str (REPO_ROOT ), "/root/ml-models" , copy = True )
43+ .add_local_file (
44+ "/Users/mulgogi/src/interscript/rababa/sadeed_evaluator.py" ,
45+ "/opt/rababa/sadeed_evaluator.py" ,
46+ copy = True ,
47+ )
48+ .add_local_dir (
49+ "/Users/mulgogi/src/interscript/rababa/data/sadeed-diac-25" ,
50+ "/opt/rababa/data/sadeed-diac-25" ,
51+ copy = True ,
52+ )
3953 .workdir ("/root/ml-models" )
4054)
4155
6882 "mode" : "sequence" , # cross-tokenizer: teacher generates, student trains CE
6983 "note" : "umt5 (sentencepiece) teacher -> ByT5-small byte student; +5pp PER gate" ,
7084 },
85+ "ara-diac-small" : {
86+ # r5 paragraph-context teacher (2.68 DER-CE windowed @1400B,
87+ # RELEASE-FROZEN) -> ByT5-small student. Contract decode is
88+ # GREEDY with generation cap 2x window (eval_sadeed_windowed).
89+ # Corpus: r5-units joined paragraph units (src = stripped
90+ # diacritics, teacher regenerates the labels).
91+ "teacher" : "rababa_arabic_byt5/run-005-context/best" ,
92+ "teacher_volume" : "rababa" ,
93+ "student_init" : "google/byt5-small" ,
94+ "train" : "r5-units/domain.txt" ,
95+ "train_extra" : ["r5-units/replay.txt" ],
96+ "unit_limits" : [24000 , 6000 ],
97+ "max_len" : 1450 ,
98+ "label_beams" : "1" ,
99+ "out" : "rababa_arabic_distill_small/run-002" ,
100+ "mode" : "sequence" ,
101+ "note" : "gate <= teacher_der + 0.5pp windowed DER-CE (prompt target 3.18 from 2.68)" ,
102+ },
71103 "fas-g2p-small" : {
72104 "teacher" : "persian_g2p/run-001/best" ,
73105 "teacher_volume" : "persian" ,
@@ -505,49 +537,75 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict:
505537 student .train ()
506538
507539 class Pairs (Dataset ):
508- def __init__ (self , paths : Path | list [Path ], max_len : int = 384 ):
509- if isinstance (paths , Path ):
510- paths = [paths ]
540+ def __init__ (self , files : list [tuple [Path , int ]], max_len : int = 1450 ):
541+ # jsonl files carry {src, tgt} rows; .txt unit files are
542+ # single-column diacritized paragraph units (the r5 corpus):
543+ # src = diacritics stripped, tgt = the unit itself, capped at
544+ # max_len bytes, seeded shuffle then per-file limit
545+ import random
546+ import re
547+
511548 self .rows = []
512549 seen = set ()
513- for path in paths :
514- for line in path .read_text (encoding = "utf-8" ).splitlines ():
515- if not line .strip ():
516- continue
517- try :
518- row = json .loads (line )
519- except json .JSONDecodeError :
520- continue
521- s = (row .get ("src" ) or "" ).strip ()
522- if s and s not in seen and len (s .encode ()) <= max_len :
523- seen .add (s )
524- self .rows .append ((s , (row .get ("tgt" ) or "" ).strip ()))
550+ for path , limit in files :
551+ if path .suffix == ".jsonl" :
552+ for line in path .read_text (encoding = "utf-8" ).splitlines ():
553+ if not line .strip ():
554+ continue
555+ try :
556+ row = json .loads (line )
557+ except json .JSONDecodeError :
558+ continue
559+ s = (row .get ("src" ) or "" ).strip ()
560+ if s and s not in seen and len (s .encode ()) <= 384 :
561+ seen .add (s )
562+ self .rows .append ((s , (row .get ("tgt" ) or "" ).strip ()))
563+ else :
564+ diac = re .compile ("[ً-ٰٟۖ-ۭ]" )
565+ units = [
566+ u .strip ()
567+ for u in path .read_text (encoding = "utf-8" , errors = "ignore" ).splitlines ()
568+ if u .strip ()
569+ ]
570+ random .Random (42 ).shuffle (units )
571+ for unit in units [:limit ]:
572+ if len (unit .encode ()) > max_len :
573+ continue
574+ src = diac .sub ("" , unit ).strip ()
575+ if src and src not in seen :
576+ seen .add (src )
577+ self .rows .append ((src , unit ))
525578
526579 def __len__ (self ):
527580 return len (self .rows )
528581
529582 def __getitem__ (self , i ):
530583 return self .rows [i ]
531584
585+ train_cap = int (spec .get ("max_len" , 384 ))
586+
532587 def collate (batch ):
533588 # byte-level tokens: a 2,000-char Wikipedia sentence is 2,000
534589 # tokens — without truncation a single long pair OOMs the A10G
535590 src = student_tok (
536591 [s for s , _ in batch ], padding = True , truncation = True ,
537- max_length = 384 , return_tensors = "pt" ,
592+ max_length = train_cap , return_tensors = "pt" ,
538593 )
539594 labels = student_tok (
540595 [t for _ , t in batch ], padding = True , truncation = True ,
541- max_length = 384 , return_tensors = "pt" ,
596+ max_length = train_cap , return_tensors = "pt" ,
542597 ).input_ids
543598 labels [labels == student_tok .pad_token_id ] = - 100
544599 return src .input_ids , src .attention_mask , labels
545600
546- train_files = [train_path ] + [
547- Path (data_vol ) / p for p in spec .get ("train_extra" , [])
548- ]
601+ unit_limits = [int (x ) for x in spec .get ("unit_limits" , [0 ])]
602+ train_files = [(train_path , unit_limits [0 ] if unit_limits else 0 )]
603+ for i , p in enumerate (spec .get ("train_extra" , [])):
604+ lim = unit_limits [i + 1 ] if i + 1 < len (unit_limits ) else 0
605+ train_files .append ((Path (data_vol ) / p , lim ))
549606 train_ds = Pairs (train_files )
550607 print (f"[{ spec_id } ] train pairs: { len (train_ds )} from { len (train_files )} files" , flush = True )
608+ label_beams = int (spec .get ("label_beams" , 4 ))
551609
552610 # Step 1: teacher generates labels (beam-4) for the full corpus.
553611 # Resumable: evictions mid-labeling are routine on long jobs —
@@ -570,10 +628,14 @@ def collate(batch):
570628 if todo :
571629 print (f"[{ spec_id } ] labeling { len (todo )} remaining..." , flush = True )
572630
573- def label_batch (batch , max_len : int = 384 ):
631+ seq_max = int (spec .get ("max_len" , 384 ))
632+
633+ def label_batch (batch , max_len : int = 0 ):
574634 # lone-src OOM fallback truncates once, then skips: never
575635 # recurse on the same shape (torch 2.x renames the OOM
576636 # exception class, so match by message)
637+ if not max_len :
638+ max_len = seq_max
577639 try :
578640 enc = teacher_tok (
579641 [s for s , _ in batch ],
@@ -584,7 +646,9 @@ def label_batch(batch, max_len: int = 384):
584646 ).to ("cuda" )
585647 with torch .inference_mode ():
586648 out = teacher .generate (
587- ** enc , max_new_tokens = max_len , num_beams = 4
649+ # r5 contract: generation cap = 2x window bytes
650+ # (diacritized output runs 1.4-1.6x input)
651+ ** enc , max_new_tokens = 2 * max_len , num_beams = label_beams
588652 )
589653 return [decode_joined (teacher_tok , o ) for o in out ]
590654 except RuntimeError as e :
@@ -604,7 +668,7 @@ def label_batch(batch, max_len: int = 384):
604668 # deterministic token-budget batching: sort by length so long
605669 # srcs land in small batches — no OOM roulette
606670 todo .sort (key = lambda p : len (p [0 ].encode ()))
607- budget = 16 * 200
671+ budget = 32 * max ( 200 , seq_max )
608672 batches : list [list [tuple [str , str ]]] = []
609673 cur : list [tuple [str , str ]] = []
610674 cur_max = 0
@@ -746,6 +810,133 @@ def eval_main(spec: str = "heb-diac-small", limit: int = 0) -> None:
746810 print (evaluate .remote (spec , limit ))
747811
748812
813+ @app .function (
814+ gpu = "A10G" ,
815+ cpu = 8 ,
816+ memory = 32 * 1024 ,
817+ timeout = 5 * 3600 ,
818+ volumes = {
819+ "/datasets" : DATASETS ,
820+ "/checkpoints" : CHECKPOINTS ,
821+ "/secryst-checkpoints" : SECRYST_CHECKPOINTS ,
822+ "/secryst-datasets" : SECRYST_DATASETS ,
823+ "/persian-checkpoints" : PERSIAN_CHECKPOINTS ,
824+ },
825+ )
826+ def evaluate_der (spec_id : str , window : int = 1400 , limit : int = 0 ) -> dict :
827+ """Windowed SadeedDiac-25 DER-CE of teacher vs student, replicating
828+ rababa eval_sadeed_windowed.py at the r5 window (1400B): strip
829+ diacritics, split at word boundaries, greedy decode with 2x window
830+ cap, stitch, project haraqat onto the input letters (zero-skip),
831+ DER-CE via the Misraj evaluator. Gate: teacher + 0.5pp
832+ (DISTILL-SOURCE-PROMPT: 3.18 target from the 2.68 teacher)."""
833+ import difflib
834+ import re
835+ from pathlib import Path
836+
837+ import pyarrow .parquet as pq
838+ import torch
839+ from transformers import AutoModelForSeq2SeqLM , AutoTokenizer
840+
841+ spec = SPECS [spec_id ]
842+ vol_map = {
843+ "rababa" : "/checkpoints" ,
844+ "secryst" : "/secryst-checkpoints" ,
845+ "persian" : "/persian-checkpoints" ,
846+ }
847+ teacher_path = (spec ["teacher" ] if spec .get ("teacher_is_hub" )
848+ else str (Path (vol_map [spec .get ("teacher_volume" , "rababa" )]) / spec ["teacher" ]))
849+ student_path = Path (vol_map [spec .get ("teacher_volume" , "rababa" )]) / spec ["out" ] / "best"
850+
851+ tok = AutoTokenizer .from_pretrained ("google/byt5-small" )
852+ teacher = AutoModelForSeq2SeqLM .from_pretrained (teacher_path ).to ("cuda" ).eval ()
853+ student = AutoModelForSeq2SeqLM .from_pretrained (str (student_path )).to ("cuda" ).eval ()
854+
855+ diac = re .compile ("[ً-ٰٟۖ-ۭ]" )
856+ table = pq .read_table ("/opt/rababa/data/sadeed-diac-25/train.parquet" )
857+ inputs = [diac .sub ("" , t ) for t in table .column ("input" ).to_pylist ()]
858+ gts = table .column ("output" ).to_pylist ()
859+ if limit :
860+ inputs , gts = inputs [:limit ], gts [:limit ]
861+
862+ def split_windows (text : str ) -> list [str ]:
863+ if len (text .encode ()) <= window :
864+ return [text ]
865+ wins , cur , n = [], [], 0
866+ for w in text .split ():
867+ c = len (w .encode ()) + 1
868+ if cur and n + c > window :
869+ wins .append (" " .join (cur ))
870+ cur , n = [], 0
871+ cur .append (w )
872+ n += c
873+ if cur :
874+ wins .append (" " .join (cur ))
875+ return wins
876+
877+ def project_haraqat (pred : str , text : str ) -> str :
878+ haraqat = ["" ]
879+ for ch in pred :
880+ if diac .match (ch ):
881+ haraqat [- 1 ] += ch
882+ else :
883+ haraqat .append ("" )
884+ haraqat = haraqat [1 :]
885+ pred_letters = [c for c in pred if not diac .match (c )]
886+ text_letters = [c for c in text if not diac .match (c )]
887+ sm = difflib .SequenceMatcher (None , text_letters , pred_letters , autojunk = False )
888+ out = []
889+ for op , i1 , i2 , j1 , _ in sm .get_opcodes ():
890+ if op == "equal" :
891+ for k in range (i2 - i1 ):
892+ out .append (text_letters [i1 + k ] + haraqat [j1 + k ])
893+ else :
894+ for k in range (i1 , i2 ):
895+ out .append (text_letters [k ])
896+ return "" .join (out )
897+
898+ def der_ce (model ) -> dict :
899+ windows , counts = [], []
900+ for text in inputs :
901+ ws = split_windows (text )
902+ counts .append (len (ws ))
903+ windows .extend (ws )
904+ preds = []
905+ with torch .no_grad ():
906+ for i in range (0 , len (windows ), 8 ):
907+ batch = windows [i : i + 8 ]
908+ enc = tok (batch , return_tensors = "pt" , padding = True , truncation = True ,
909+ max_length = window ).to ("cuda" )
910+ with torch .autocast ("cuda" , torch .bfloat16 ):
911+ gen = model .generate (** enc , max_new_tokens = window * 2 , num_beams = 1 )
912+ preds .extend (tok .batch_decode (gen , skip_special_tokens = True ))
913+ k = 0
914+ paragraphs = []
915+ for text , c in zip (inputs , counts , strict = True ):
916+ paragraphs .append (project_haraqat (" " .join (preds [k : k + c ]), text ))
917+ k += c
918+
919+ import sys
920+
921+ sys .path .insert (0 , "/opt/rababa" )
922+ from sadeed_evaluator import ArabicDiacritizationEvaluator as E
923+
924+ _ , _ , total_der , _ , _ = E .caculate_errors_on_sentences (
925+ paragraphs , gts , gt_missing_diacritic_is_error = False
926+ )
927+ return {"der_ce" : round (100 * total_der , 4 ), "n" : len (inputs )}
928+
929+ result = {"teacher" : der_ce (teacher ), "student" : der_ce (student )}
930+ result ["gate_delta" ] = round (result ["student" ]["der_ce" ] - result ["teacher" ]["der_ce" ], 4 )
931+ result ["gate_pass" ] = result ["gate_delta" ] <= 0.5
932+ return result
933+
934+
749935@app .local_entrypoint ()
750936def eval_per (spec : str = "tha-g2p-small" , limit : int = 0 ) -> None :
751937 print (evaluate_per .remote (spec , limit ))
938+
939+
940+ @app .local_entrypoint ()
941+ def eval_der (spec : str = "ara-diac-small" , limit : int = 0 ) -> None :
942+ print (evaluate_der .remote (spec , limit ))
0 commit comments