| |
| """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_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 = cfg.get('training', {}) |
| |
| 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() |
|
|
