File size: 4,048 Bytes
f211247 | 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 | #!/usr/bin/env python3
"""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)
# Data
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, # pretrain epochs (shortened for 1-day target)
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'))
# Also save symlink/copy as checkpoint.pt for seq2seq_ft
import shutil
# Save in trainer format; train_image loads via HF from_pretrained
log(f"DONE: {output}/checkpoint")
if __name__ == '__main__':
main()
|