File size: 4,510 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""ByT5 Hittite transliteration fine-tune.
Input: cuneiform Unicode (or transliteration tokens)
Output: transliteration string
Data: TLHdig 280K Hittite parallel pairs
Sequence packing (block-diagonal attention) için HF native DataCollator kullanıyoruz.
"""
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/transliteration_hitit_only.yaml'))
ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/hitit_byt5/'))
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('name', 'google/byt5-small')
log(f"Loading model {model_name}...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)
# Data: hittite_parallel (cuneiform → transliteration)
data_path = ROOT / 'datasets/processed/tlhdig_corpora/hittite_parallel.jsonl'
log(f"Loading data: {data_path}")
examples = []
with open(data_path) 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 examples: {len(examples)}")
val_frac = cfg.get('data', {}).get('val_fraction', 0.1)
n_val = int(len(examples) * val_frac)
val_examples = examples[:n_val]
train_examples = examples[n_val:]
log(f"Train: {len(train_examples)}, Val: {len(val_examples)}")
def preprocess(batch):
inputs = tokenizer(batch['input_text'], max_length=512, truncation=True, padding=False)
targets = tokenizer(batch['target_text'], max_length=256, truncation=True, padding=False)
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
training = cfg.get('training', {})
# BF16 availability check (torchrun sometimes loses this)
import torch
bf16_supported = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
use_bf16 = training.get('bf16', True) and bf16_supported
targs = Seq2SeqTrainingArguments(
output_dir=str(output),
num_train_epochs=training.get('epochs', 30),
per_device_train_batch_size=training.get('batch_size', 32),
per_device_eval_batch_size=training.get('batch_size', 32),
learning_rate=training.get('lr', 5e-5),
weight_decay=training.get('weight_decay', 0.01),
warmup_steps=training.get('warmup_steps', 1000),
bf16=use_bf16,
fp16=False,
logging_steps=100,
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,
gradient_accumulation_steps=1,
dataloader_num_workers=4,
report_to='none',
)
collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding='longest')
trainer = Seq2SeqTrainer(
model=model,
args=targs,
train_dataset=train_ds,
eval_dataset=val_ds,
data_collator=collator,
tokenizer=tokenizer,
)
log("Starting training...")
trainer.train()
log("Saving final model...")
trainer.save_model(str(output / 'best'))
tokenizer.save_pretrained(str(output / 'best'))
log(f"DONE: {output}")
if __name__ == '__main__':
main()
|