| |
| """Image → Transliteration seq2seq model. |
| |
| Reference: SumTablets Simmons 2024, HATFormer 2024. |
| Architecture: ViT encoder + ByT5 decoder + cross-attention. |
| """ |
| import torch |
| import torch.nn as nn |
| from pathlib import Path |
|
|
| class VisionByT5(nn.Module): |
| """ViT encoder + ByT5 decoder for cuneiform transliteration.""" |
| |
| def __init__(self, encoder_name="facebook/dinov2-large", |
| decoder_name="google/byt5-small", |
| image_size=384): |
| super().__init__() |
| |
| |
| from transformers import AutoModel, T5ForConditionalGeneration, AutoTokenizer |
| self.encoder = AutoModel.from_pretrained(encoder_name) |
| self.decoder = T5ForConditionalGeneration.from_pretrained(decoder_name) |
| self.tokenizer = AutoTokenizer.from_pretrained(decoder_name) |
| |
| |
| enc_dim = self.encoder.config.hidden_size |
| dec_dim = self.decoder.config.d_model |
| self.enc_to_dec_proj = nn.Linear(enc_dim, dec_dim) |
| |
| |
| self._apply_lora(self.encoder, r=32, alpha=64) |
| |
| def _apply_lora(self, module, r=32, alpha=64): |
| try: |
| from peft import LoraConfig, get_peft_model |
| config = LoraConfig( |
| r=r, lora_alpha=alpha, |
| target_modules=["query", "key", "value", "dense"], |
| lora_dropout=0.1, bias="none" |
| ) |
| module = get_peft_model(module, config) |
| except ImportError: |
| print("peft not installed; LoRA skipped", flush=True) |
| |
| def forward(self, pixel_values, labels=None): |
| """ |
| pixel_values: (B, 3, H, W) |
| labels: (B, T) — target byte sequence |
| """ |
| |
| enc_out = self.encoder(pixel_values=pixel_values) |
| enc_hidden = enc_out.last_hidden_state |
| |
| |
| enc_hidden = self.enc_to_dec_proj(enc_hidden) |
| |
| |
| outputs = self.decoder( |
| encoder_outputs=(enc_hidden,), |
| labels=labels, |
| ) |
| return outputs |
| |
| @torch.no_grad() |
| def generate(self, pixel_values, max_length=512, num_beams=5): |
| enc_out = self.encoder(pixel_values=pixel_values) |
| enc_hidden = self.enc_to_dec_proj(enc_out.last_hidden_state) |
| |
| return self.decoder.generate( |
| encoder_outputs=(enc_hidden,), |
| max_length=max_length, |
| num_beams=num_beams, |
| length_penalty=0.6, |
| no_repeat_ngram_size=3, |
| early_stopping=True, |
| ) |
|
|
| if __name__ == '__main__': |
| print("Vision-ByT5 seq2seq model architecture template hazır.") |
| print("Training data: tablet image + transliteration pair") |
| print("Expected chrF (SumTablets baseline): 0.9755") |
|
|
