{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":121150,"databundleVersionId":14976537,"sourceType":"competition"},{"sourceId":88039,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":73936,"modelId":98760},{"sourceId":690442,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":523415,"modelId":537419}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nEnhanced Old Assyrian Translation using NLLB-200\nOptimized Strategy: Smart augmentation of train.csv only\n(OCR data is bibliographic, not suitable for extraction)\n\"\"\"\n\nimport os, re, math, random\nimport numpy as np\nimport pandas as pd\nfrom typing import List, Tuple, Dict\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n)\n\n# ============================================================================\n# CONFIGURATION\n# ============================================================================\n\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(SEED)\n\n# Paths\nTRAIN_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\nTEST_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\nLEXICON_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv\"\n\n# Model paths\nMODEL_PATHS = [\n    \"/kaggle/input/nllb-distilled-600-million/transformers/default/1/NLLB 200 distilled 600m\",\n    \"/kaggle/input/nllb-200-600m\",\n    \"facebook/nllb-200-distilled-600M\",\n]\n\n# NLLB language codes\nSRC_LANG = \"akk_Latn\"\nTGT_LANG = \"eng_Latn\"\n\n# Training hyperparameters (optimized for augmented data)\nMAX_SRC_LEN = 160\nMAX_TGT_LEN = 120\nBATCH_SIZE = 2\nGRAD_ACCUM = 8\nLEARNING_RATE = 1.5e-5\nEPOCHS = 30\nWARMUP_RATIO = 0.05\nWEIGHT_DECAY = 0.01\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Device: {DEVICE}\")\nif torch.cuda.is_available():\n    print(f\"GPUs: {torch.cuda.device_count()}\")\n\n# ============================================================================\n# TEXT PREPROCESSING\n# ============================================================================\n\nUNICODE_MAP = {\n    \"ā\": \"aa\", \"Ā\": \"AA\", \"ē\": \"ee\", \"Ē\": \"EE\",\n    \"ī\": \"ii\", \"Ī\": \"II\", \"ū\": \"uu\", \"Ū\": \"UU\",\n    \"á\": \"a2\", \"à\": \"a3\", \"é\": \"e2\", \"è\": \"e3\",\n    \"í\": \"i2\", \"ì\": \"i3\", \"ú\": \"u2\", \"ù\": \"u3\",\n    \"š\": \"sh\", \"Š\": \"SH\", \"ṣ\": \"s2\", \"Ṣ\": \"S2\",\n    \"ṭ\": \"t2\", \"Ṭ\": \"T2\", \"ḫ\": \"h2\", \"Ḫ\": \"H2\",\n}\n\nSUB_MAP = str.maketrans({\n    \"₀\": \"0\", \"₁\": \"1\", \"₂\": \"2\", \"₃\": \"3\", \"₄\": \"4\",\n    \"₅\": \"5\", \"₆\": \"6\", \"₇\": \"7\", \"₈\": \"8\", \"₉\": \"9\"\n})\n\ndef clean_ws(x: str) -> str:\n    if pd.isna(x):\n        return \"\"\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\ndef normalize_akkadian(text: str) -> str:\n    if pd.isna(text):\n        return \"\"\n    \n    text = str(text).strip()\n    text = text.translate(SUB_MAP)\n    \n    for orig, repl in UNICODE_MAP.items():\n        text = text.replace(orig, repl)\n    \n    text = re.sub(r\"\\{([^}]+)\\}\", r\" DET_\\1 \", text)\n    text = text.replace(\"…\", \" [GAP] \")\n    text = re.sub(r\"\\[x\\]\", \" [GAP] \", text)\n    text = re.sub(r\"\\b[xX]{1,3}\\b\", \" [GAP] \", text)\n    text = re.sub(r\"[˹˺]\", \"\", text)\n    text = re.sub(r\"[!?/:]\", \" \", text)\n    text = re.sub(r\"[ʾʿˀˁ']\", \"\", text)\n    text = re.sub(r\"\\[([^\\]]+)\\]\", r\" \\1 \", text)\n    text = re.sub(r\"<([^>]+)>\", r\" \\1 \", text)\n    text = text.replace(\".\", \" . \")\n    \n    return clean_ws(text)\n\ndef normalize_english(text: str) -> str:\n    text = clean_ws(text)\n    text = text.replace(\"\"\", '\"').replace(\"\"\", '\"')\n    text = text.replace(\"'\", \"'\").replace(\"'\", \"'\")\n    return text\n\n# ============================================================================\n# LEXICON\n# ============================================================================\n\ndef load_lexicon_map(lexicon_path: str) -> Dict[str, str]:\n    \"\"\"Load lexicon for morphology-aware augmentation\"\"\"\n    try:\n        lex_df = pd.read_csv(lexicon_path)\n        word_map = {}\n        \n        for _, row in lex_df.iterrows():\n            if pd.notna(row['form']) and pd.notna(row['lexeme']):\n                form = str(row['form']).strip().lower()\n                lexeme = str(row['lexeme']).strip().lower()\n                \n                if form and lexeme and len(form) >= 2:\n                    word_map[form] = lexeme\n                    \n                    if pd.notna(row['norm']):\n                        norm = str(row['norm']).strip().lower()\n                        if norm:\n                            word_map[norm] = lexeme\n        \n        print(f\"✓ Loaded {len(word_map)} lexicon mappings\")\n        return word_map\n    except Exception as e:\n        print(f\"⚠ Could not load lexicon: {e}\")\n        return {}\n\n# ============================================================================\n# SENTENCE ALIGNMENT\n# ============================================================================\n\ndef split_by_lines(text: str) -> List[str]:\n    if pd.isna(text):\n        return []\n    text = str(text)\n    parts = [p.strip() for p in text.split('\\n') if p.strip()]\n    if len(parts) <= 1:\n        parts = re.split(r'\\s{2,}|;\\s+', text)\n        parts = [p.strip() for p in parts if p.strip()]\n    return parts\n\ndef split_sentences(text: str) -> List[str]:\n    text = clean_ws(text)\n    if not text:\n        return []\n    sents = re.split(r'(?<=[.!?])\\s+', text)\n    return [s.strip() for s in sents if s.strip()]\n\ndef smart_align(src_parts: List[str], tgt_parts: List[str], max_merge: int = 2) -> List[Tuple[str, str]]:\n    \"\"\"Smart alignment using length ratios\"\"\"\n    def get_length(s: str) -> int:\n        return max(1, len(s.split()))\n    \n    i = j = 0\n    pairs = []\n    \n    while i < len(src_parts) and j < len(tgt_parts):\n        best_score = float('inf')\n        best_config = None\n        \n        for si in range(1, min(max_merge + 1, len(src_parts) - i + 1)):\n            for tj in range(1, min(max_merge + 1, len(tgt_parts) - j + 1)):\n                src_chunk = \" \".join(src_parts[i:i+si])\n                tgt_chunk = \" \".join(tgt_parts[j:j+tj])\n                \n                ls = get_length(src_chunk)\n                lt = get_length(tgt_chunk)\n                \n                score = abs(math.log(max(ls, 1) / max(lt, 1)))\n                score += 0.1 * (si + tj)\n                \n                if score < best_score:\n                    best_score = score\n                    best_config = (si, tj, src_chunk, tgt_chunk)\n        \n        if best_config:\n            si, tj, src_chunk, tgt_chunk = best_config\n            pairs.append((src_chunk, tgt_chunk))\n            i += si\n            j += tj\n        else:\n            break\n    \n    # Append remainders to last pair\n    if pairs:\n        if i < len(src_parts):\n            remainder = \" \".join(src_parts[i:])\n            pairs[-1] = (pairs[-1][0] + \" \" + remainder, pairs[-1][1])\n        if j < len(tgt_parts):\n            remainder = \" \".join(tgt_parts[j:])\n            pairs[-1] = (pairs[-1][0], pairs[-1][1] + \" \" + remainder)\n    \n    return pairs\n\n# ============================================================================\n# AUGMENTATION STRATEGIES\n# ============================================================================\n\ndef build_augmented_training_data(train_df: pd.DataFrame, lexicon_map: Dict[str, str]) -> pd.DataFrame:\n    \"\"\"\n    Build comprehensive training data with intelligent augmentation\n    \"\"\"\n    print(\"\\n\" + \"=\"*80)\n    print(\"BUILDING AUGMENTED TRAINING DATA\")\n    print(\"=\"*80)\n    \n    # Step 1: Extract base sentence pairs\n    print(\"\\n[Step 1/3] Extracting sentence-level pairs...\")\n    all_pairs = []\n    \n    for idx, row in train_df.iterrows():\n        src_raw = row[\"transliteration\"]\n        tgt_raw = row[\"translation\"]\n        \n        src_parts = split_by_lines(src_raw)\n        tgt_parts = split_sentences(tgt_raw)\n        \n        if not src_parts or not tgt_parts:\n            continue\n        \n        doc_pairs = smart_align(src_parts, tgt_parts, max_merge=2)\n        \n        for src, tgt in doc_pairs:\n            src_norm = normalize_akkadian(src)\n            tgt_norm = normalize_english(tgt)\n            \n            if src_norm and tgt_norm and len(src_norm) > 10 and len(tgt_norm) > 15:\n                all_pairs.append({\n                    \"src\": src_norm,\n                    \"tgt\": tgt_norm,\n                    \"aug_type\": \"original\",\n                    \"doc_id\": idx\n                })\n    \n    base_df = pd.DataFrame(all_pairs)\n    print(f\"✓ Created {len(base_df)} base pairs\")\n    \n    # Step 2: Apply augmentation strategies\n    print(\"\\n[Step 2/3] Applying augmentation strategies...\")\n    augmented = []\n    \n    for _, row in base_df.iterrows():\n        src = row['src']\n        tgt = row['tgt']\n        doc_id = row['doc_id']\n        \n        # Always add original\n        augmented.append(row.to_dict())\n        \n        # Aug 1: Remove [GAP] markers (critical for test set)\n        if '[GAP]' in src:\n            src_no_gaps = re.sub(r'\\[GAP\\]', '', src)\n            src_no_gaps = re.sub(r'\\s+', ' ', src_no_gaps).strip()\n            if len(src_no_gaps) > 10 and src_no_gaps != src:\n                augmented.append({\n                    'src': src_no_gaps,\n                    'tgt': tgt,\n                    'aug_type': 'no_gaps',\n                    'doc_id': doc_id\n                })\n        \n        # Aug 2: Simplify DET_ markers (DET_d -> d)\n        if 'DET_' in src:\n            src_simple = re.sub(r'DET_(\\w+)', r'\\1', src)\n            if src_simple != src and len(src_simple) > 10:\n                augmented.append({\n                    'src': src_simple,\n                    'tgt': tgt,\n                    'aug_type': 'simple_det',\n                    'doc_id': doc_id\n                })\n        \n        # Aug 3: Lemmatize using lexicon (morphology normalization)\n        if lexicon_map:\n            src_words = src.split()\n            lemmatized_words = []\n            changes = 0\n            \n            for word in src_words:\n                word_clean = word.strip('.,;:')\n                word_lower = word_clean.lower()\n                \n                if word_lower in lexicon_map:\n                    lemma = lexicon_map[word_lower]\n                    if lemma != word_lower:\n                        lemmatized_words.append(lemma)\n                        changes += 1\n                    else:\n                        lemmatized_words.append(word)\n                else:\n                    lemmatized_words.append(word)\n            \n            # Add if we made 1-3 meaningful changes\n            if 1 <= changes <= 3:\n                src_lemmatized = ' '.join(lemmatized_words)\n                if src_lemmatized != src and len(src_lemmatized) > 10:\n                    augmented.append({\n                        'src': src_lemmatized,\n                        'tgt': tgt,\n                        'aug_type': 'lemmatized',\n                        'doc_id': doc_id\n                    })\n        \n        # Aug 4: Combined (no gaps + simple det)\n        if '[GAP]' in src and 'DET_' in src:\n            src_combined = re.sub(r'\\[GAP\\]', '', src)\n            src_combined = re.sub(r'DET_(\\w+)', r'\\1', src_combined)\n            src_combined = re.sub(r'\\s+', ' ', src_combined).strip()\n            if len(src_combined) > 10 and src_combined != src:\n                augmented.append({\n                    'src': src_combined,\n                    'tgt': tgt,\n                    'aug_type': 'combined',\n                    'doc_id': doc_id\n                })\n    \n    aug_df = pd.DataFrame(augmented)\n    \n    # Remove exact duplicates\n    aug_df = aug_df.drop_duplicates(subset=['src', 'tgt'])\n    \n    print(f\"✓ Total after augmentation: {len(aug_df)} pairs\")\n    print(f\"\\n  Breakdown by type:\")\n    for aug_type in sorted(aug_df['aug_type'].unique()):\n        count = len(aug_df[aug_df['aug_type'] == aug_type])\n        pct = 100 * count / len(aug_df)\n        print(f\"    - {aug_type:12s}: {count:4d} ({pct:5.1f}%)\")\n    \n    # Step 3: Create diverse batches\n    print(\"\\n[Step 3/3] Creating diverse training batches...\")\n    \n    # Shuffle within each document, then interleave\n    doc_groups = []\n    for doc_id in aug_df['doc_id'].unique():\n        doc_pairs = aug_df[aug_df['doc_id'] == doc_id].copy()\n        doc_pairs = doc_pairs.sample(frac=1.0, random_state=SEED).reset_index(drop=True)\n        doc_groups.append(doc_pairs)\n    \n    # Interleave documents for diverse batches\n    final_df = pd.concat(doc_groups, ignore_index=True)\n    final_df = final_df.sample(frac=1.0, random_state=SEED).reset_index(drop=True)\n    \n    print(f\"✓ Created diverse batch ordering\")\n    print(f\"  Unique documents: {final_df['doc_id'].nunique()}\")\n    \n    return final_df\n\n# ============================================================================\n# MODEL & DATASET\n# ============================================================================\n\ndef load_model_tokenizer():\n    for path in MODEL_PATHS:\n        try:\n            print(f\"Trying: {path}\")\n            tokenizer = AutoTokenizer.from_pretrained(path)\n            model = AutoModelForSeq2SeqLM.from_pretrained(path)\n            is_nllb = hasattr(tokenizer, 'lang_code_to_id') or 'nllb' in path.lower()\n            print(f\"✓ Loaded: {path}\")\n            print(f\"  Model type: {'NLLB-200' if is_nllb else 'Generic Seq2Seq'}\")\n            return tokenizer, model, is_nllb, path\n        except Exception as e:\n            print(f\"  Failed: {e}\")\n            continue\n    raise RuntimeError(\"Could not load any model!\")\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, tokenizer, is_nllb: bool):\n        self.src = df[\"src\"].tolist()\n        self.tgt = df[\"tgt\"].tolist()\n        self.tokenizer = tokenizer\n        self.is_nllb = is_nllb\n        \n        if self.is_nllb:\n            if hasattr(self.tokenizer, 'src_lang'):\n                self.tokenizer.src_lang = SRC_LANG\n            if hasattr(self.tokenizer, 'tgt_lang'):\n                self.tokenizer.tgt_lang = TGT_LANG\n    \n    def __len__(self):\n        return len(self.src)\n    \n    def __getitem__(self, idx):\n        src_text = self.src[idx]\n        tgt_text = self.tgt[idx]\n        \n        if self.is_nllb:\n            if hasattr(self.tokenizer, 'src_lang'):\n                self.tokenizer.src_lang = SRC_LANG\n            if hasattr(self.tokenizer, 'tgt_lang'):\n                self.tokenizer.tgt_lang = TGT_LANG\n        \n        inputs = self.tokenizer(\n            src_text,\n            max_length=MAX_SRC_LEN,\n            truncation=True,\n            padding=False,\n        )\n        \n        labels = self.tokenizer(\n            text_target=tgt_text,\n            max_length=MAX_TGT_LEN,\n            truncation=True,\n            padding=False,\n        )\n        \n        inputs[\"labels\"] = labels[\"input_ids\"]\n        return inputs\n\n# ============================================================================\n# MAIN\n# ============================================================================\n\ndef main():\n    print(\"=\" * 80)\n    print(\"ENHANCED OLD ASSYRIAN TRANSLATION - NLLB-200\")\n    print(\"Optimized Strategy: Smart Augmentation of train.csv\")\n    print(\"=\" * 80)\n    \n    # Load data\n    print(\"\\n[1/6] Loading data...\")\n    train_df = pd.read_csv(TRAIN_PATH)\n    test_df = pd.read_csv(TEST_PATH)\n    print(f\"✓ Train: {len(train_df)} documents\")\n    print(f\"✓ Test: {len(test_df)} sentences\")\n    \n    # Load lexicon\n    print(\"\\n[2/6] Loading lexicon...\")\n    lexicon_map = load_lexicon_map(LEXICON_PATH)\n    \n    # Build augmented training data\n    print(\"\\n[3/6] Building augmented training data...\")\n    all_pairs = build_augmented_training_data(train_df, lexicon_map)\n    \n    print(f\"\\n{'='*80}\")\n    print(f\"✓ FINAL TRAINING DATA\")\n    print(f\"{'='*80}\")\n    print(f\"  Total pairs: {len(all_pairs)}\")\n    print(f\"  Avg src length: {all_pairs['src'].str.split().str.len().mean():.1f} tokens\")\n    print(f\"  Avg tgt length: {all_pairs['tgt'].str.split().str.len().mean():.1f} tokens\")\n    print(f\"  Unique documents: {all_pairs['doc_id'].nunique()}\")\n    print(f\"  Training epochs: {EPOCHS}\")\n    print(f\"  Total training examples: {len(all_pairs) * EPOCHS:,}\")\n    print(f\"{'='*80}\")\n    \n    # Load model\n    print(\"\\n[4/6] Loading model...\")\n    tokenizer, model, is_nllb, model_path = load_model_tokenizer()\n    model.to(DEVICE)\n    \n    if hasattr(model, 'gradient_checkpointing_enable'):\n        model.gradient_checkpointing_enable()\n        print(\"✓ Gradient checkpointing enabled\")\n    \n    # Prepare datasets\n    print(\"\\n[5/6] Preparing datasets...\")\n    all_pairs = all_pairs.sample(frac=1.0, random_state=SEED).reset_index(drop=True)\n    \n    # Small validation set (3%)\n    n_val = max(100, int(len(all_pairs) * 0.03))\n    val_df = all_pairs.iloc[:n_val].reset_index(drop=True)\n    trn_df = all_pairs.iloc[n_val:].reset_index(drop=True)\n    \n    print(f\"✓ Train: {len(trn_df)} pairs\")\n    print(f\"✓ Val: {len(val_df)} pairs\")\n    \n    train_ds = AkkadianDataset(trn_df, tokenizer, is_nllb)\n    val_ds = AkkadianDataset(val_df, tokenizer, is_nllb)\n    \n    data_collator = DataCollatorForSeq2Seq(\n        tokenizer=tokenizer,\n        model=model,\n        padding=True,\n    )\n    \n    # Clear cache\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n        import gc\n        gc.collect()\n    \n    # Training setup\n    print(\"\\n[6/6] Training setup...\")\n    total_steps = (len(train_ds) // (BATCH_SIZE * GRAD_ACCUM)) * EPOCHS\n    warmup_steps = int(total_steps * WARMUP_RATIO)\n    \n    print(f\"  Effective batch size: {BATCH_SIZE * GRAD_ACCUM * torch.cuda.device_count()}\")\n    print(f\"  Total steps: {total_steps}\")\n    print(f\"  Warmup steps: {warmup_steps}\")\n    print(f\"  Epochs: {EPOCHS}\")\n    \n    training_args = Seq2SeqTrainingArguments(\n        output_dir=\"./akkadian_nllb\",\n        eval_strategy=\"steps\",\n        eval_steps=max(50, total_steps // 20),\n        save_steps=999999,\n        save_total_limit=1,\n        logging_steps=50,\n        \n        learning_rate=LEARNING_RATE,\n        per_device_train_batch_size=BATCH_SIZE,\n        per_device_eval_batch_size=BATCH_SIZE,\n        gradient_accumulation_steps=GRAD_ACCUM,\n        \n        num_train_epochs=EPOCHS,\n        warmup_steps=warmup_steps,\n        weight_decay=WEIGHT_DECAY,\n        \n        gradient_checkpointing=True,\n        gradient_checkpointing_kwargs={\"use_reentrant\": False},\n        \n        predict_with_generate=True,\n        generation_max_length=MAX_TGT_LEN,\n        generation_num_beams=4,\n        \n        fp16=True if DEVICE == \"cuda\" else False,\n        dataloader_num_workers=0,\n        dataloader_pin_memory=False,\n        \n        save_only_model=True,\n        load_best_model_at_end=False,\n        metric_for_best_model=\"eval_loss\",\n        \n        report_to=\"none\",\n        remove_unused_columns=False,\n        \n        optim=\"adafactor\",\n        max_grad_norm=0.5,\n    )\n    \n    trainer = Seq2SeqTrainer(\n        model=model,\n        args=training_args,\n        train_dataset=train_ds,\n        eval_dataset=val_ds,\n        data_collator=data_collator,\n        processing_class=tokenizer,\n    )\n    \n    # Train\n    print(\"\\n\" + \"=\"*80)\n    print(\"STARTING TRAINING\")\n    print(\"=\"*80)\n    \n    trainer.train()\n    \n    print(\"\\n✓ Training complete!\")\n    \n    # Save model\n    print(\"\\nSaving model...\")\n    try:\n        model.save_pretrained(\"./final_model\", safe_serialization=True)\n        tokenizer.save_pretrained(\"./final_model\")\n        print(\"✓ Model saved to ./final_model/\")\n    except Exception as e:\n        print(f\"⚠ Could not save model: {e}\")\n        torch.save(model.state_dict(), \"./model_weights.pt\")\n        print(\"✓ Weights saved to ./model_weights.pt\")\n    \n    # Inference\n    print(\"\\n\" + \"=\"*80)\n    print(\"INFERENCE ON TEST SET\")\n    print(\"=\"*80)\n    \n    model.eval()\n    \n    def finalize_translation(text: str) -> str:\n        \"\"\"Post-process translation\"\"\"\n        text = clean_ws(text)\n        if not text:\n            return \"...\"\n        \n        # Remove leading seal/envelope markers\n        if re.match(r\"^\\s*(Seal of|Sealed by|Envelope|Copy of)\", text, re.I):\n            parts = re.split(r'(?<=[.!?])\\s+', text)\n            if len(parts) >= 2:\n                text = \" \".join(parts[1:])\n        \n        # Split into sentences and clean\n        parts = re.split(r'(?<=[.!?])\\s+', text)\n        parts = [re.sub(r'[.!?]+\\s*$', \"\", p).strip() for p in parts if p.strip()]\n        \n        if not parts:\n            return \"...\"\n        \n        # Join with semicolons, end with period\n        result = \"; \".join(parts)\n        if not result.endswith(\".\"):\n            result += \".\"\n        \n        return result\n    \n    def translate(src_text: str) -> str:\n        \"\"\"Translate a single source text\"\"\"\n        src_norm = normalize_akkadian(src_text)\n        \n        inputs = tokenizer(\n            src_norm,\n            return_tensors=\"pt\",\n            max_length=MAX_SRC_LEN,\n            truncation=True,\n        ).to(DEVICE)\n        \n        with torch.no_grad():\n            gen_kwargs = {\n                \"num_beams\": 5,\n                \"max_new_tokens\": MAX_TGT_LEN,\n                \"length_penalty\": 1.0,\n                \"early_stopping\": True,\n            }\n            \n            if is_nllb and hasattr(tokenizer, 'lang_code_to_id'):\n                tgt_lang_id = tokenizer.lang_code_to_id.get(TGT_LANG)\n                if tgt_lang_id:\n                    gen_kwargs[\"forced_bos_token_id\"] = tgt_lang_id\n            \n            outputs = model.generate(**inputs, **gen_kwargs)\n        \n        pred = tokenizer.decode(outputs[0], skip_special_tokens=True)\n        return finalize_translation(pred)\n    \n    print(f\"\\nTranslating {len(test_df)} test examples...\")\n    predictions = []\n    \n    for i, text in enumerate(test_df[\"transliteration\"].tolist()):\n        pred = translate(text)\n        predictions.append(pred)\n        \n        if (i + 1) % 100 == 0 or i + 1 == len(test_df):\n            print(f\"  Progress: {i+1}/{len(test_df)}\")\n    \n    # Save submission\n    submission = pd.DataFrame({\n        \"id\": test_df[\"id\"],\n        \"translation\": predictions\n    })\n    \n    submission.to_csv(\"submission.csv\", index=False)\n    print(\"\\n✓ Submission saved to submission.csv\")\n    \n    # Show sample predictions\n    print(\"\\n\" + \"=\"*80)\n    print(\"SAMPLE PREDICTIONS\")\n    print(\"=\"*80)\n    \n    for i in range(min(3, len(test_df))):\n        src_text = test_df.iloc[i]['transliteration']\n        pred_text = predictions[i]\n        \n        print(f\"\\n[Example {i+1}]\")\n        print(f\"Source: {src_text[:100]}{'...' if len(src_text) > 100 else ''}\")\n        print(f\"Translation: {pred_text[:150]}{'...' if len(pred_text) > 150 else ''}\")\n    \n    print(\"\\n\" + \"=\"*80)\n    print(\"TRAINING COMPLETE!\")\n    print(f\"Total training pairs: {len(all_pairs)}\")\n    print(f\"Epochs: {EPOCHS}\")\n    print(f\"Augmentation types: {', '.join(sorted(all_pairs['aug_type'].unique()))}\")\n    print(\"=\"*80)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T10:25:16.870484Z","iopub.execute_input":"2025-12-19T10:25:16.871041Z","execution_failed":"2025-12-19T10:26:40.622Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}