| |
| """ByT5 text-only pretrain on Hittite parallel corpus. |
| |
| Cuneiform Unicode → Transliteration, 280K pairs. |
| Reference: SumTablets ACL ML4AL 2024. |
| """ |
| import os, sys, json, argparse, time |
| from pathlib import Path |
| import yaml |
|
|
| ROOT = Path("/arf/scratch/stakan/hitit-proje") |
|
|
| def log(msg): print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/seq2seq_hitit.yaml')) |
| ap.add_argument('--phase', default='pretrain') |
| ap.add_argument('--data', default=str(ROOT / 'datasets/processed/tlhdig_corpora/hittite_parallel.jsonl')) |
| ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/seq2seq_text_pretrain/')) |
| args = ap.parse_args() |
|
|
| cfg = yaml.safe_load(open(args.config)) |
| output = Path(args.output) |
| output.mkdir(parents=True, exist_ok=True) |
|
|
| try: |
| import torch |
| from transformers import (AutoTokenizer, T5ForConditionalGeneration, |
| Seq2SeqTrainer, Seq2SeqTrainingArguments, |
| DataCollatorForSeq2Seq) |
| from datasets import Dataset |
| except ImportError as e: |
| log(f"Missing library: {e}"); sys.exit(1) |
|
|
| model_name = cfg.get('model', {}).get('decoder', {}).get('pretrained', 'google/byt5-small') |
| log(f"Phase: {args.phase}, Model: {model_name}") |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = T5ForConditionalGeneration.from_pretrained(model_name) |
|
|
| |
| log(f"Loading {args.data}") |
| examples = [] |
| with open(args.data) as f: |
| for line in f: |
| r = json.loads(line) |
| cune = r.get('cuneiform') or '' |
| translit = r.get('transliteration') or '' |
| if cune and translit and len(cune) < 200 and len(translit) < 400: |
| examples.append({'input_text': cune, 'target_text': translit}) |
| log(f"Total: {len(examples)} pairs") |
|
|
| n_val = int(len(examples) * 0.05) |
| val_examples = examples[:n_val] |
| train_examples = examples[n_val:] |
|
|
| def preprocess(batch): |
| inputs = tokenizer(batch['input_text'], max_length=512, truncation=True) |
| targets = tokenizer(batch['target_text'], max_length=512, truncation=True) |
| inputs['labels'] = targets['input_ids'] |
| return inputs |
|
|
| train_ds = Dataset.from_list(train_examples).map(preprocess, batched=True, |
| remove_columns=['input_text', 'target_text']) |
| val_ds = Dataset.from_list(val_examples).map(preprocess, batched=True, |
| remove_columns=['input_text', 'target_text']) |
|
|
| training = cfg.get('training', {}) |
| targs = Seq2SeqTrainingArguments( |
| output_dir=str(output), |
| num_train_epochs=3, |
| per_device_train_batch_size=training.get('batch_size', 8), |
| per_device_eval_batch_size=16, |
| learning_rate=1e-4, |
| weight_decay=0.01, |
| warmup_steps=2000, |
| bf16=True, |
| fp16=False, |
| logging_steps=200, |
| eval_strategy='epoch', |
| save_strategy='epoch', |
| save_total_limit=2, |
| load_best_model_at_end=True, |
| metric_for_best_model='eval_loss', |
| greater_is_better=False, |
| dataloader_num_workers=4, |
| report_to='none', |
| ) |
|
|
| trainer = Seq2SeqTrainer( |
| model=model, args=targs, |
| train_dataset=train_ds, eval_dataset=val_ds, |
| data_collator=DataCollatorForSeq2Seq(tokenizer, model=model), |
| tokenizer=tokenizer, |
| ) |
|
|
| log("Training text-only pretrain...") |
| trainer.train() |
|
|
| trainer.save_model(str(output / 'checkpoint')) |
| tokenizer.save_pretrained(str(output / 'checkpoint')) |
| |
| import shutil |
| |
| log(f"DONE: {output}/checkpoint") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
