# --- # jupyter: # jupytext: # text_representation: # extension: .py # format_name: percent # kernelspec: # display_name: Python 3 # language: python # name: python3 # --- # %% [markdown] # # T5-Base Akkadian Translation Training - Enriched V2 (Stable) # # Training with enriched dataset (3,031 pairs) - STABLE VERSION # # Changes from v1: # - Lower learning rate (5e-5 vs 3e-4) to prevent gradient explosion # - Disabled fp16 to avoid overflow with Unicode diacritics # - More warmup steps # - Uses v3 dataset with fixed fractions # # Uses: manwithacat/akkadian-enriched-data (v3) # %% [markdown] # ## Configuration # %% CONFIG = { # Model - T5-Base (223M params) "model_name": "google-t5/t5-base", "max_src_len": 512, "max_tgt_len": 192, # Training - STABLE settings "num_epochs": 25, # Slightly fewer epochs "learning_rate": 5e-5, # Much lower - standard for T5 fine-tuning "batch_size": 8, "gradient_accumulation": 2, "warmup_ratio": 0.15, # More warmup "weight_decay": 0.01, "max_grad_norm": 0.5, # Stricter clipping # Checkpointing "save_steps": 200, "eval_steps": 200, "logging_steps": 50, "save_total_limit": 3, # Data "val_fraction": 0.05, "seed": 42, } # %% import os import re import random import numpy as np import pandas as pd from pathlib import Path import torch from torch.utils.data import Dataset from transformers import ( AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer, ) # Set seeds SEED = CONFIG["seed"] random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED) # Paths - Kaggle environment COMP_DIR = Path("/kaggle/input/deep-past-initiative-machine-translation") ENRICHED_DIR = Path("/kaggle/input/akkadian-enriched-data") OUTPUT_DIR = Path("/kaggle/working/t5-akkadian-enriched-v2") OUTPUT_DIR.mkdir(exist_ok=True, parents=True) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" print(f"Device: {DEVICE}") if DEVICE == "cuda": print(f"GPU: {torch.cuda.get_device_name(0)}") print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") # %% [markdown] # ## Character Normalization # %% CHAR_NORMALIZE_MAP = { 'á': 'a₂', 'à': 'a₃', 'é': 'e₂', 'è': 'e₃', 'í': 'i₂', 'ì': 'i₃', 'ú': 'u₂', 'ù': 'u₃', 'ṣ': 'ṣ', 'Ṣ': 'Ṣ', 'ṭ': 'ṭ', 'Ṭ': 'Ṭ', 'ḫ': 'ḫ', 'Ḫ': 'Ḫ', ''': 'ʾ', "'": 'ʾ', ''': 'ʾ', } MULTI_CHAR_PATTERNS = [('sz', 'š'), ('SZ', 'Š')] def normalize_transliteration(text: str) -> str: """Normalize transliteration characters to canonical Unicode form.""" if pd.isna(text): return "" text = str(text) for pattern, replacement in MULTI_CHAR_PATTERNS: text = text.replace(pattern, replacement) result = [] for char in text: result.append(CHAR_NORMALIZE_MAP.get(char, char)) return ''.join(result) # %% [markdown] # ## Preprocessing Functions # %% def preprocess_source(text, prefix="translate Akkadian to English: "): """Preprocess source transliteration for T5.""" text = str(text) # Normalize character variants text = normalize_transliteration(text) # Clean editorial markup but keep structure text = re.sub(r'[!?„"""\'\[\]<>˹˺…+]', '', text) text = text.replace('/', ' ') # Normalize spacing text = ' '.join(text.split()) return prefix + text # %% [markdown] # ## Load Enriched Data (v3 with fixed fractions) # %% # Load enriched training data enriched_df = pd.read_csv(ENRICHED_DIR / "train_enriched.csv") print(f"Enriched training data: {len(enriched_df)} pairs") # Check for fraction artifacts import re frac_check = enriched_df['transliteration'].str.contains(r'[⅓⅔½]', regex=True, na=False).sum() print(f"Rows with Unicode fractions: {frac_check}") # Load test data for validation test_df = pd.read_csv(COMP_DIR / "test.csv") print(f"Competition test samples: {len(test_df)}") # %% [markdown] # ## Prepare Training Pairs # %% # Prepare pairs with preprocessing pairs_df = pd.DataFrame({ 'src': enriched_df['transliteration'].apply(preprocess_source), 'tgt': enriched_df['translation'], }) # Filter empty pairs_df = pairs_df[ (pairs_df['src'].str.len() > 0) & (pairs_df['tgt'].str.len() > 0) ].drop_duplicates().reset_index(drop=True) print(f"Training pairs after filtering: {len(pairs_df)}") # Show samples print("\nSample pairs:") for i in range(min(3, len(pairs_df))): print(f"\n[{i}] SRC: {pairs_df.iloc[i]['src'][:80]}...") print(f" TGT: {pairs_df.iloc[i]['tgt'][:80]}...") # %% [markdown] # ## Dataset Class # %% class MTDataset(Dataset): def __init__(self, df, tokenizer, max_src, max_tgt): self.src = df["src"].tolist() self.tgt = df["tgt"].tolist() self.tokenizer = tokenizer self.max_src = max_src self.max_tgt = max_tgt def __len__(self): return len(self.src) def __getitem__(self, idx): x = self.tokenizer(self.src[idx], max_length=self.max_src, truncation=True) y = self.tokenizer(text_target=self.tgt[idx], max_length=self.max_tgt, truncation=True) x["labels"] = y["input_ids"] return x # %% [markdown] # ## Load Model # %% print(f"\nLoading model: {CONFIG['model_name']}...") try: tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_name"]) model = AutoModelForSeq2SeqLM.from_pretrained(CONFIG["model_name"]) except Exception as e: print(f"Failed to load {CONFIG['model_name']}, trying t5-base...") tokenizer = AutoTokenizer.from_pretrained("t5-base") model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") model.to(DEVICE) print(f"Model: {model.config.model_type}, Parameters: {sum(p.numel() for p in model.parameters()):,}") data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model) # %% [markdown] # ## Prepare Datasets # %% # Split data pairs_df = pairs_df.sample(frac=1.0, random_state=SEED).reset_index(drop=True) n_val = max(1, int(len(pairs_df) * CONFIG["val_fraction"])) val_df = pairs_df.iloc[:n_val].reset_index(drop=True) trn_df = pairs_df.iloc[n_val:].reset_index(drop=True) print(f"Train: {len(trn_df)}, Val: {len(val_df)}") train_ds = MTDataset(trn_df, tokenizer, CONFIG["max_src_len"], CONFIG["max_tgt_len"]) val_ds = MTDataset(val_df, tokenizer, CONFIG["max_src_len"], CONFIG["max_tgt_len"]) # %% [markdown] # ## Training - STABLE SETTINGS # %% effective_batch = CONFIG["batch_size"] if DEVICE == "cuda" else 2 effective_accum = CONFIG["gradient_accumulation"] if DEVICE == "cuda" else 1 training_args = Seq2SeqTrainingArguments( output_dir=str(OUTPUT_DIR / "checkpoints"), eval_strategy="steps", eval_steps=CONFIG["eval_steps"], save_steps=CONFIG["save_steps"], save_total_limit=CONFIG["save_total_limit"], logging_steps=CONFIG["logging_steps"], learning_rate=CONFIG["learning_rate"], per_device_train_batch_size=effective_batch, per_device_eval_batch_size=effective_batch, gradient_accumulation_steps=effective_accum, num_train_epochs=CONFIG["num_epochs"], warmup_ratio=CONFIG["warmup_ratio"], weight_decay=CONFIG["weight_decay"], max_grad_norm=CONFIG["max_grad_norm"], predict_with_generate=True, generation_max_length=CONFIG["max_tgt_len"], # STABILITY: Disable fp16 to avoid overflow with Unicode fp16=False, bf16=False, # Would use if available, but P100 doesn't support it report_to="none", remove_unused_columns=False, load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, # Additional stability dataloader_pin_memory=True, optim="adamw_torch", ) trainer = Seq2SeqTrainer( model=model, args=training_args, train_dataset=train_ds, eval_dataset=val_ds, data_collator=data_collator, ) print("\n" + "="*60) print("Starting STABLE Enriched Training V2") print("="*60) print(f"Epochs: {CONFIG['num_epochs']}") print(f"Learning rate: {CONFIG['learning_rate']} (lower for stability)") print(f"fp16: DISABLED (to avoid overflow)") print(f"Max grad norm: {CONFIG['max_grad_norm']} (stricter clipping)") print(f"Training pairs: {len(trn_df)}") print(f"Effective batch size: {effective_batch * effective_accum}") print("="*60) trainer.train() print("Training complete!") # %% [markdown] # ## Save Model # %% final_model_path = OUTPUT_DIR / "final" final_model_path.mkdir(exist_ok=True, parents=True) print(f"\nSaving model to {final_model_path}...") model.save_pretrained(final_model_path) tokenizer.save_pretrained(final_model_path) import json with open(final_model_path / "training_config.json", "w") as f: json.dump(CONFIG, f, indent=2) # Model metadata for Kaggle model_meta = { "ownerSlug": "manwithacat", "modelSlug": "akkadian-t5-enriched-v2", "title": "Akkadian T5 Enriched V2 (Stable)", "subtitle": "T5-Base trained with stable hyperparameters", "isPrivate": False, "description": f"""T5-Base fine-tuned for Akkadian translation with STABLE training. ## Changes from V1 - Learning rate: 5e-5 (was 3e-4) - fp16: Disabled (prevents overflow with Unicode) - Stricter gradient clipping: 0.5 (was 1.0) - More warmup: 15% (was 10%) ## Training Data - {len(enriched_df)} pairs (enriched dataset v3) - Fixed fraction normalization (⅓ not 0⅓) ## Training - Epochs: {CONFIG['num_epochs']} - Learning rate: {CONFIG['learning_rate']} - Batch size: {effective_batch * effective_accum} """ } with open(final_model_path / "model-metadata.json", "w") as f: json.dump(model_meta, f, indent=2) print("Model saved successfully!") # %% [markdown] # ## Validation # %% def generate_translation(src_text: str) -> str: """Generate translation using trained model.""" src_preprocessed = preprocess_source(src_text) enc = tokenizer( src_preprocessed, return_tensors="pt", max_length=CONFIG["max_src_len"], truncation=True, ).to(DEVICE) with torch.no_grad(): out_ids = model.generate( **enc, num_beams=4, max_new_tokens=CONFIG["max_tgt_len"], length_penalty=1.0, repetition_penalty=1.5, # Less aggressive no_repeat_ngram_size=3, early_stopping=True, ) pred = tokenizer.decode(out_ids[0], skip_special_tokens=True) return pred model.eval() print("\n" + "="*60) print("Validation Samples") print("="*60) for i in range(min(4, len(test_df))): row = test_df.iloc[i] trans = generate_translation(row["transliteration"]) print(f"\nTest ID {row['id']}:") print(f" Akkadian: {row['transliteration'][:60]}...") print(f" Translation: {trans[:100]}...") print("\n" + "="*60) print("Stable Enriched V2 Training complete!") print("="*60)