{"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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14188526,"sourceType":"datasetVersion","datasetId":9046477},{"sourceId":14239027,"sourceType":"datasetVersion","datasetId":9084458},{"sourceId":14243325,"sourceType":"datasetVersion","datasetId":9087187},{"sourceId":14248946,"sourceType":"datasetVersion","datasetId":9091299},{"sourceId":14255176,"sourceType":"datasetVersion","datasetId":9095663},{"sourceId":166632804,"sourceType":"kernelVersion"},{"sourceId":11371,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":5171,"modelId":3533},{"sourceId":208042,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":5171,"modelId":3533}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"print(\"\\n... PIP INSTALLS ...\\n\")\n\n#!pip install -q -U --upgrade pip\n#!pip install -q -U --upgrade tensorflow-hub\n#!pip install -q -U --upgrade tensorflow-cpu \n#!pip install -q -U --upgrade tf-keras\n#!pip install -q -U --upgrade keras-nlp \n#!pip install -q -U --upgrade keras>=3\n!pip install --no-index /kaggle/input/llm-wheels/tensorflow_hub-0.16.1-py2.py3-none-any.whl --find-links=/kaggle/input/llm-wheels/\n!pip install --no-index /kaggle/input/llm-wheels/tensorflow_cpu-2.16.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl --find-links=/kaggle/input/llm-wheels/\n!pip install --no-index /kaggle/input/llm-wheels/tf_keras-2.16.0-py3-none-any.whl --find-links=/kaggle/input/llm-wheels/\n!pip install --no-index /kaggle/input/llm-wheels/keras_nlp-0.8.2-py3-none-any.whl --find-links=/kaggle/input/llm-wheels/\n!pip install --no-index /kaggle/input/llm-wheels/keras-3.0.5-py3-none-any.whl --find-links=/kaggle/input/llm-wheels/\n\nuse_7b = False #False:Gemma 2B; True: Gemma 7B\n\nprint(\"\\n... IMPORTS ...\\n\")\nimport os\n\n\nprint(\"\\n... ENVIRONMENT SETUP (LOGGING/KERAS-BACKEND) ...\\n\")\n# Set backend\n# Pre-allocate 90% of TPU memory to minimize memory fragmentation and allocation overhead\n# Disable logging/warning that may bloat the output\n\nkeras_backend = \"jax\"\nallocation_fraction = 0.9\ntf_min_log_level = '3'\nos.environ[\"KERAS_BACKEND\"] = keras_backend\nos.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = str(allocation_fraction)\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = str(tf_min_log_level)\n\n# Finish imports\nimport jax\nimport keras\nimport keras_nlp\n\n# Other imports\nimport time\nimport copy\n\nprint(jax.devices())\nprint(keras.backend.backend())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:50:20.354927Z","iopub.execute_input":"2025-12-22T10:50:20.355300Z","iopub.status.idle":"2025-12-22T10:50:53.151861Z","shell.execute_reply.started":"2025-12-22T10:50:20.355262Z","shell.execute_reply":"2025-12-22T10:50:53.151021Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import jax\nimport keras\nimport keras_nlp\n\n# Note we use jax.devices() instead of keras.distribution.list_devices()\ndef initialize_device_mesh(\n    shape: tuple[int, int] | None = None,\n    batch_axis_name: str = \"batch\",\n    model_axis_name: str = \"model\",\n    device_type: str = \"gpu\",  # \"gpu\" for GPUs, \"tpu\" for TPUs, \"cpu\" for CPU\n) -> keras.distribution.DeviceMesh:\n    \"\"\"Initializes and returns a DeviceMesh for distributing computation across devices.\n\n    For GPU: uses jax.devices(\"gpu\"). If no GPUs are available, falls back to jax.devices().\n    If `shape` is None, it auto-chooses:\n      - (1, N) when N>1 devices (model axis across N GPUs)\n      - (1, 1) when only 1 device\n    \"\"\"\n    # Prefer the requested device type if present (e.g., \"gpu\").\n    devices = jax.devices(device_type)\n    if not devices:\n        # Fallback (e.g., if no GPU is present, you still get CPU devices)\n        devices = jax.devices()\n\n    n = len(devices)\n\n    if shape is None:\n        # Default: shard model axis across all devices if possible\n        shape = (1, n) if n > 1 else (1, 1)\n    else:\n        # Basic sanity check: shape product must match device count\n        prod = shape[0] * shape[1]\n        if prod != n:\n            raise ValueError(\n                f\"DeviceMesh shape {shape} has {prod} slots, but found {n} devices ({device_type}). \"\n                f\"Use shape=(1,{n}) or set shape to match {n}.\"\n            )\n\n    return keras.distribution.DeviceMesh(\n        shape=shape,\n        axis_names=[batch_axis_name, model_axis_name],\n        devices=devices,\n    )\n\n\ndef configure_layout_map(\n    device_mesh: keras.distribution.DeviceMesh,\n    model_axis_name: str = \"model\",\n) -> keras.distribution.LayoutMap:\n    \"\"\"Configures and returns a LayoutMap for model weight distribution.\n\n    If the model axis has size 1 (e.g., single GPU), we do NOT apply sharding rules.\n    \"\"\"\n    layout_map = keras.distribution.LayoutMap(device_mesh)\n\n    # If there's only 1 device along the model axis, sharding doesn't help (and can error).\n    model_axis_size = device_mesh.shape[device_mesh.axis_names.index(model_axis_name)]\n    if model_axis_size <= 1:\n        return layout_map  # everything replicated\n\n    # Shard selected weights across the model axis\n    layout_map[\"token_embedding/embeddings\"] = (None, model_axis_name)\n    layout_map[\"decoder_block.*attention.*(query|key|value).*kernel\"] = (None, model_axis_name, None)\n    layout_map[\"decoder_block.*attention_output.*kernel\"] = (None, None, model_axis_name)\n    layout_map[\"decoder_block.*ffw_gating.*kernel\"] = (model_axis_name, None)\n    layout_map[\"decoder_block.*ffw_linear.*kernel\"] = (None, model_axis_name)\n\n    return layout_map\n\n\ndef update_distribution_strategy(\n    device_mesh: keras.distribution.DeviceMesh,\n    layout_map: keras.distribution.LayoutMap,\n    batch_axis_name: str = \"batch\",\n    model_axis_name: str = \"model\",\n) -> None:\n    \"\"\"Updates the keras backend distribution strategy.\n\n    Uses ModelParallel only when there are >1 devices along the model axis.\n    Otherwise leaves default (single device) behavior.\n    \"\"\"\n    model_axis_size = device_mesh.shape[device_mesh.axis_names.index(model_axis_name)]\n    if model_axis_size <= 1:\n        # Single GPU (or single device): no distribution needed.\n        # If you previously set a distribution strategy in this process,\n        # you can restart the kernel/session for a clean slate.\n        return\n\n    model_parallel = keras.distribution.ModelParallel(\n        device_mesh, layout_map, batch_dim_name=batch_axis_name\n    )\n    keras.distribution.set_distribution(model_parallel)\n\n\ndef get_distributed_gemma(model_preset_or_path: str) -> keras_nlp.models.GemmaCausalLM:\n    \"\"\"Loads Gemma from a preset name or a local preset path.\"\"\"\n    return keras_nlp.models.GemmaCausalLM.from_preset(model_preset_or_path)\n\n\ndef do_gemma_prep(\n    return_device_mesh: bool = True,\n    return_layout_map: bool = True,\n    device_type: str = \"gpu\",\n    shape: tuple[int, int] | None = None,\n):\n    \"\"\"Prepares device mesh + layout map + distribution strategy for GPU.\"\"\"\n    device_mesh = initialize_device_mesh(shape=shape, device_type=device_type)\n    layout_map = configure_layout_map(device_mesh)\n    update_distribution_strategy(device_mesh, layout_map)\n\n    out = {}\n    if return_device_mesh:\n        out[\"device_mesh\"] = device_mesh\n    if return_layout_map:\n        out[\"layout_map\"] = layout_map\n    return out\n\n\n# ---- Usage ----\nsetup_objects = do_gemma_prep(device_type=\"gpu\")  # auto-detect GPU count, auto-shape\n\nuse_7b = False\nuse_v3 = False\n\nif use_v3:\n    model_path = \"/kaggle/input/gemma/keras/gemma_2b_en/3\"\nelse:\n    model_path = \"/kaggle/input/gemma/keras/gemma_2b_en/2\"\n\nmodel = get_distributed_gemma(model_path)\nGLOBAL_TOKENIZER = keras_nlp.models.GemmaTokenizer.from_preset(model_path)\n\nprint(f\"SETUP ARTIFACTS:\\n{setup_objects}\\n\\nMODEL OBJECT:\\n{model}\")\nprint(\"\\n\\n\\n... MODEL SUMMARY ...\\n\\n\")\nmodel.summary()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:50:53.154281Z","iopub.execute_input":"2025-12-22T10:50:53.155434Z","iopub.status.idle":"2025-12-22T10:51:37.035246Z","shell.execute_reply.started":"2025-12-22T10:50:53.155383Z","shell.execute_reply":"2025-12-22T10:51:37.034671Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, re, math, random\nimport numpy as np\nimport pandas as pd\n\nimport torch\nfrom torch.utils.data import Dataset\n\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:51:37.036145Z","iopub.execute_input":"2025-12-22T10:51:37.036446Z","iopub.status.idle":"2025-12-22T10:51:51.361504Z","shell.execute_reply.started":"2025-12-22T10:51:37.036412Z","shell.execute_reply":"2025-12-22T10:51:51.360923Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\nSEED = 42\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\n\nTRAIN_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\nTEST_PATH  = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n\n# Put your local model path here if you added it as a Kaggle Dataset:\n# e.g. \"/kaggle/input/byt5-small\" or \"/kaggle/input/mt5-small\"\nMODEL_PATH_CANDIDATES = [\"/kaggle/input/byt5-small\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:51:51.362360Z","iopub.execute_input":"2025-12-22T10:51:51.362930Z","iopub.status.idle":"2025-12-22T10:51:51.372778Z","shell.execute_reply.started":"2025-12-22T10:51:51.362904Z","shell.execute_reply":"2025-12-22T10:51:51.372159Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)\n\ntrain_df = pd.read_csv(TRAIN_PATH)\ntest_df  = pd.read_csv(TEST_PATH)\n\nprint(\"train:\", train_df.shape, \"test:\", test_df.shape)\n\n\ndef clean_ws(x: str) -> str:\n    if pd.isna(x): return \"\"\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\ndef split_src_transliteration(doc: str):\n    \"\"\"Split transliteration doc into 'lines' as proxy for sentences.\"\"\"\n    doc = \"\" if pd.isna(doc) else str(doc)\n    # Often transliterations are line-based with \\n\n    parts = re.split(r\"\\n+\", doc)\n    parts = [clean_ws(p) for p in parts if clean_ws(p)]\n    # fallback: if no newlines, split on \" ; \" or long spaces\n    if len(parts) <= 1:\n        parts = re.split(r\"\\s{2,}|;\\s+\", clean_ws(doc))\n        parts = [clean_ws(p) for p in parts if clean_ws(p)]\n    return parts\n\ndef split_tgt_english(doc: str):\n    \"\"\"Split English translation into sentences (simple regex).\"\"\"\n    doc = clean_ws(doc)\n    if not doc:\n        return []\n    # Keep punctuation as boundary\n    sents = re.split(r\"(?<=[.!?])\\s+\", doc)\n    sents = [clean_ws(s) for s in sents if clean_ws(s)]\n    return sents\n\ndef tok_len(s: str) -> int:\n    # length proxy for alignment; byte-level later anyway\n    return max(1, len(s.split()))\n\ndef align_monotone(src_units, tgt_units, max_src_merge=2, max_tgt_merge=2):\n    \"\"\"\n    Monotone greedy alignment by matching length proportions.\n    Produces pairs (src_text, tgt_text), allowing merges 1-2 units.\n    \"\"\"\n    i = j = 0\n    pairs = []\n    while i < len(src_units) and j < len(tgt_units):\n        best = None\n        # try combinations: 1..max merge\n        for si in range(1, max_src_merge+1):\n            for tj in range(1, max_tgt_merge+1):\n                if i+si <= len(src_units) and j+tj <= len(tgt_units):\n                    src_chunk = \" \".join(src_units[i:i+si])\n                    tgt_chunk = \" \".join(tgt_units[j:j+tj])\n                    # objective: minimize normalized length mismatch\n                    ls, lt = tok_len(src_chunk), tok_len(tgt_chunk)\n                    score = abs(math.log(ls/lt))\n                    if (best is None) or (score < best[0]):\n                        best = (score, si, tj, src_chunk, tgt_chunk)\n        if best is None:\n            break\n        _, si, tj, s_chunk, t_chunk = best\n        pairs.append((s_chunk, t_chunk))\n        i += si\n        j += tj\n\n    # If leftovers exist, attach them to last pair (simple & safe)\n    if pairs:\n        if i < len(src_units):\n            pairs[-1] = (pairs[-1][0] + \" \" + \" \".join(src_units[i:]), pairs[-1][1])\n        if j < len(tgt_units):\n            pairs[-1] = (pairs[-1][0], pairs[-1][1] + \" \" + \" \".join(tgt_units[j:]))\n\n    return pairs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:51:51.373969Z","iopub.execute_input":"2025-12-22T10:51:51.374235Z","iopub.status.idle":"2025-12-22T10:51:51.447643Z","shell.execute_reply.started":"2025-12-22T10:51:51.374212Z","shell.execute_reply":"2025-12-22T10:51:51.447027Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, re, math, random\nimport numpy as np\nimport pandas as pd\nfrom collections import Counter, defaultdict\nimport json\nimport pickle\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\nSEED = 1\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\nTRAIN_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\nTEST_PATH  = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)\n\ntrain_df = pd.read_csv(TRAIN_PATH)\ntest_df  = pd.read_csv(TEST_PATH)\n\nprint(\"train:\", train_df.shape, \"test:\", test_df.shape)\n\n\n# ============================================================================\n# BPE TOKENIZER IMPLEMENTATION\n# ============================================================================\n\nclass SimpleBPETokenizer:\n    \"\"\"\n    Standalone Byte Pair Encoding tokenizer for transliteration tasks.\n    No dependencies on transformers library.\n    \"\"\"\n    \n    def __init__(self, vocab_size=8000):\n        self.vocab_size = vocab_size\n        self.token_to_id = {}\n        self.id_to_token = {}\n        self.merges = []\n        \n        # Special tokens\n        self.pad_token = \"<PAD>\"\n        self.unk_token = \"<UNK>\"\n        self.bos_token = \"<BOS>\"\n        self.eos_token = \"<EOS>\"\n        \n        self.pad_token_id = 0\n        self.unk_token_id = 1\n        self.bos_token_id = 2\n        self.eos_token_id = 3\n        \n        # Initialize with special tokens\n        self._initialize_vocab()\n        \n    def _initialize_vocab(self):\n        \"\"\"Initialize vocabulary with special tokens.\"\"\"\n        self.token_to_id = {\n            self.pad_token: self.pad_token_id,\n            self.unk_token: self.unk_token_id,\n            self.bos_token: self.bos_token_id,\n            self.eos_token: self.eos_token_id,\n        }\n        self.id_to_token = {v: k for k, v in self.token_to_id.items()}\n    \n    def _get_word_freqs(self, texts):\n        \"\"\"Get word frequencies from texts.\"\"\"\n        word_freqs = Counter()\n        for text in texts:\n            if pd.isna(text):\n                continue\n            text = str(text).strip()\n            if not text:\n                continue\n            # Split on whitespace\n            words = text.split()\n            for word in words:\n                # Represent word as sequence of characters + end marker\n                word_freqs[tuple(word) + ('</w>',)] += 1\n        return word_freqs\n    \n    def _get_pair_freqs(self, word_freqs):\n        \"\"\"Get frequencies of adjacent character pairs.\"\"\"\n        pair_freqs = Counter()\n        for word, freq in word_freqs.items():\n            for i in range(len(word) - 1):\n                pair_freqs[(word[i], word[i+1])] += freq\n        return pair_freqs\n    \n    def _merge_pair(self, word, pair):\n        \"\"\"Merge all occurrences of a pair in a word.\"\"\"\n        new_word = []\n        i = 0\n        while i < len(word):\n            if i < len(word) - 1 and (word[i], word[i+1]) == pair:\n                new_word.append(word[i] + word[i+1])\n                i += 2\n            else:\n                new_word.append(word[i])\n                i += 1\n        return tuple(new_word)\n    \n    def train(self, texts, verbose=True):\n        \"\"\"\n        Train BPE on a corpus of texts.\n        \n        Args:\n            texts: List of strings to train on\n            verbose: Print training progress\n        \"\"\"\n        if verbose:\n            print(f\"\\nTraining BPE tokenizer on {len(texts)} texts...\")\n        \n        # Get initial word frequencies (character-level)\n        word_freqs = self._get_word_freqs(texts)\n        \n        if verbose:\n            print(f\"Unique words: {len(word_freqs)}\")\n        \n        # Add all characters to vocabulary\n        chars = set()\n        for word in word_freqs.keys():\n            chars.update(word)\n        \n        for char in sorted(chars):\n            if char not in self.token_to_id:\n                token_id = len(self.token_to_id)\n                self.token_to_id[char] = token_id\n                self.id_to_token[token_id] = char\n        \n        if verbose:\n            print(f\"Initial vocab size (with chars): {len(self.token_to_id)}\")\n        \n        # Perform BPE merges\n        num_merges = self.vocab_size - len(self.token_to_id)\n        \n        for merge_idx in range(num_merges):\n            # Get pair frequencies\n            pair_freqs = self._get_pair_freqs(word_freqs)\n            \n            if not pair_freqs:\n                if verbose:\n                    print(f\"No more pairs to merge at iteration {merge_idx}\")\n                break\n            \n            # Find most frequent pair\n            best_pair = max(pair_freqs, key=pair_freqs.get)\n            \n            # Merge this pair in all words\n            new_word_freqs = {}\n            for word, freq in word_freqs.items():\n                new_word = self._merge_pair(word, best_pair)\n                new_word_freqs[new_word] = freq\n            \n            word_freqs = new_word_freqs\n            \n            # Record merge\n            self.merges.append(best_pair)\n            \n            # Add merged token to vocabulary\n            merged_token = best_pair[0] + best_pair[1]\n            if merged_token not in self.token_to_id:\n                token_id = len(self.token_to_id)\n                self.token_to_id[merged_token] = token_id\n                self.id_to_token[token_id] = merged_token\n            \n            if verbose and (merge_idx + 1) % 500 == 0:\n                print(f\"Merge {merge_idx+1}/{num_merges} | \"\n                      f\"Vocab: {len(self.token_to_id)} | \"\n                      f\"Best pair: {best_pair} (freq: {pair_freqs[best_pair]})\")\n        \n        if verbose:\n            print(f\"\\nFinal vocabulary size: {len(self.token_to_id)}\")\n            print(f\"Total merges learned: {len(self.merges)}\")\n    \n    def _tokenize_word(self, word):\n        \"\"\"Apply BPE merges to tokenize a single word.\"\"\"\n        if not word:\n            return []\n        \n        # Start with character-level\n        tokens = list(word) + ['</w>']\n        \n        # Apply merges in order\n        for merge_pair in self.merges:\n            i = 0\n            new_tokens = []\n            while i < len(tokens):\n                if (i < len(tokens) - 1 and \n                    tokens[i] == merge_pair[0] and \n                    tokens[i+1] == merge_pair[1]):\n                    new_tokens.append(merge_pair[0] + merge_pair[1])\n                    i += 2\n                else:\n                    new_tokens.append(tokens[i])\n                    i += 1\n            tokens = new_tokens\n            \n            if len(tokens) == 1:\n                break\n        \n        return tokens\n    \n    def encode(self, text, add_special_tokens=True, max_length=None, \n               truncation=True, padding=False, return_tensors=None):\n        \"\"\"\n        Encode text to token IDs.\n        \n        Returns dict with 'input_ids' and 'attention_mask'.\n        \"\"\"\n        if pd.isna(text):\n            text = \"\"\n        \n        text = str(text).strip()\n        \n        # Tokenize each word\n        all_tokens = []\n        if text:\n            words = text.split()\n            for word in words:\n                word_tokens = self._tokenize_word(word)\n                all_tokens.extend(word_tokens)\n        \n        # Convert to IDs\n        token_ids = []\n        if add_special_tokens:\n            token_ids.append(self.bos_token_id)\n        \n        for token in all_tokens:\n            token_ids.append(self.token_to_id.get(token, self.unk_token_id))\n        \n        if add_special_tokens:\n            token_ids.append(self.eos_token_id)\n        \n        # Truncate if needed\n        if max_length and truncation and len(token_ids) > max_length:\n            if add_special_tokens:\n                token_ids = token_ids[:max_length-1] + [self.eos_token_id]\n            else:\n                token_ids = token_ids[:max_length]\n        \n        # Padding\n        if padding and max_length:\n            while len(token_ids) < max_length:\n                token_ids.append(self.pad_token_id)\n        \n        attention_mask = [1 if tid != self.pad_token_id else 0 for tid in token_ids]\n        \n        result = {\n            'input_ids': token_ids,\n            'attention_mask': attention_mask\n        }\n        \n        # Convert to tensors if requested\n        if return_tensors == 'pt':\n            result = {k: torch.tensor([v]) for k, v in result.items()}\n        \n        return result\n    \n    def decode(self, token_ids, skip_special_tokens=True):\n        \"\"\"Decode token IDs back to text.\"\"\"\n        if isinstance(token_ids, torch.Tensor):\n            token_ids = token_ids.tolist()\n        \n        tokens = []\n        for tid in token_ids:\n            if skip_special_tokens and tid in [\n                self.pad_token_id, self.bos_token_id, \n                self.eos_token_id, self.unk_token_id\n            ]:\n                continue\n            token = self.id_to_token.get(tid, self.unk_token)\n            tokens.append(token)\n        \n        # Join and clean up\n        text = ''.join(tokens)\n        text = text.replace('</w>', ' ')\n        return text.strip()\n    \n    def batch_decode(self, batch_ids, skip_special_tokens=True):\n        \"\"\"Decode a batch of token IDs.\"\"\"\n        return [self.decode(ids, skip_special_tokens) for ids in batch_ids]\n    \n    def __call__(self, text=None, text_target=None, **kwargs):\n        \"\"\"Make tokenizer callable.\"\"\"\n        if text is not None:\n            return self.encode(text, **kwargs)\n        elif text_target is not None:\n            return self.encode(text_target, **kwargs)\n        return {}\n    \n    def save(self, path):\n        \"\"\"Save tokenizer to file.\"\"\"\n        data = {\n            'token_to_id': self.token_to_id,\n            'id_to_token': self.id_to_token,\n            'merges': self.merges,\n            'vocab_size': self.vocab_size,\n            'special_tokens': {\n                'pad': (self.pad_token, self.pad_token_id),\n                'unk': (self.unk_token, self.unk_token_id),\n                'bos': (self.bos_token, self.bos_token_id),\n                'eos': (self.eos_token, self.eos_token_id),\n            }\n        }\n        with open(path, 'wb') as f:\n            pickle.dump(data, f)\n        print(f\"Tokenizer saved to {path}\")\n    \n    @classmethod\n    def load(cls, path):\n        \"\"\"Load tokenizer from file.\"\"\"\n        with open(path, 'rb') as f:\n            data = pickle.load(f)\n        \n        tokenizer = cls(vocab_size=data['vocab_size'])\n        tokenizer.token_to_id = data['token_to_id']\n        tokenizer.id_to_token = data['id_to_token']\n        tokenizer.merges = data['merges']\n        \n        # Restore special tokens\n        for key, (token, token_id) in data['special_tokens'].items():\n            setattr(tokenizer, f'{key}_token', token)\n            setattr(tokenizer, f'{key}_token_id', token_id)\n        \n        print(f\"Tokenizer loaded from {path}\")\n        return tokenizer\n\n\n# ============================================================================\n# DATA PREPROCESSING\n# ============================================================================\n\ndef clean_ws(x: str) -> str:\n    if pd.isna(x): return \"\"\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\ndef split_src_transliteration(doc: str):\n    doc = \"\" if pd.isna(doc) else str(doc)\n    parts = re.split(r\"\\n+\", doc)\n    parts = [clean_ws(p) for p in parts if clean_ws(p)]\n    if len(parts) <= 1:\n        parts = re.split(r\"\\s{2,}|;\\s+\", clean_ws(doc))\n        parts = [clean_ws(p) for p in parts if clean_ws(p)]\n    return parts\n\ndef split_tgt_english(doc: str):\n    doc = clean_ws(doc)\n    if not doc:\n        return []\n    sents = re.split(r\"(?<=[.!?])\\s+\", doc)\n    sents = [clean_ws(s) for s in sents if clean_ws(s)]\n    return sents\n\ndef tok_len(s: str) -> int:\n    return max(1, len(s.split()))\n\ndef align_monotone(src_units, tgt_units, max_src_merge=2, max_tgt_merge=2):\n    i = j = 0\n    pairs = []\n    while i < len(src_units) and j < len(tgt_units):\n        best = None\n        for si in range(1, max_src_merge+1):\n            for tj in range(1, max_tgt_merge+1):\n                if i+si <= len(src_units) and j+tj <= len(tgt_units):\n                    src_chunk = \" \".join(src_units[i:i+si])\n                    tgt_chunk = \" \".join(tgt_units[j:j+tj])\n                    ls, lt = tok_len(src_chunk), tok_len(tgt_chunk)\n                    score = abs(math.log(ls/lt))\n                    if (best is None) or (score < best[0]):\n                        best = (score, si, tj, src_chunk, tgt_chunk)\n        if best is None:\n            break\n        _, si, tj, s_chunk, t_chunk = best\n        pairs.append((s_chunk, t_chunk))\n        i += si\n        j += tj\n\n    if pairs:\n        if i < len(src_units):\n            pairs[-1] = (pairs[-1][0] + \" \" + \" \".join(src_units[i:]), pairs[-1][1])\n        if j < len(tgt_units):\n            pairs[-1] = (pairs[-1][0], pairs[-1][1] + \" \" + \" \".join(tgt_units[j:]))\n\n    return pairs\n\n# Create aligned sentence pairs\nprint(\"\\nCreating sentence-level alignments...\")\npairs = []\nfor idx, row in train_df.iterrows():\n    src_doc = row[\"transliteration\"]\n    tgt_doc = row[\"translation\"]\n\n    src_lines = split_src_transliteration(src_doc)\n    tgt_sents = split_tgt_english(tgt_doc)\n\n    if len(src_lines) == 0 or len(tgt_sents) == 0:\n        continue\n\n    doc_pairs = align_monotone(src_lines, tgt_sents, max_src_merge=2, max_tgt_merge=2)\n    pairs.extend(doc_pairs)\n\npairs_df = pd.DataFrame(pairs, columns=[\"src\", \"tgt\"])\npairs_df = pairs_df[(pairs_df[\"src\"].str.len() > 0) & (pairs_df[\"tgt\"].str.len() > 0)].drop_duplicates()\n\nprint(f\"Sentence-level pairs: {pairs_df.shape[0]}\")\nprint(f\"Sample pair:\")\nprint(f\"  SRC: {pairs_df['src'].iloc[0][:100]}...\")\nprint(f\"  TGT: {pairs_df['tgt'].iloc[0][:100]}...\")\n\n\n# ============================================================================\n# TRAIN BPE TOKENIZER\n# ============================================================================\n'''\nprint(\"\\n\" + \"=\"*80)\nprint(\"TRAINING BPE TOKENIZER\")\nprint(\"=\"*80)\n\n# Combine all texts for training\nall_texts = list(pairs_df[\"src\"]) + list(pairs_df[\"tgt\"])\nprint(f\"Total text samples for training: {len(all_texts)}\")\n\n# Train tokenizer\nbpe_tokenizer = SimpleBPETokenizer(vocab_size=8000)\nbpe_tokenizer.train(all_texts, verbose=True)\n\n# Save tokenizer\nbpe_tokenizer.save(\"bpe_tokenizer.pkl\")\n\n# Test tokenization\nprint(\"\\n\" + \"=\"*80)\nprint(\"TOKENIZATION EXAMPLES\")\nprint(\"=\"*80)\n\nsample_src = pairs_df[\"src\"].iloc[0]\nsample_tgt = pairs_df[\"tgt\"].iloc[0]\n\nprint(f\"\\n1. Source text:\")\nprint(f\"   {sample_src[:150]}...\")\nencoded_src = bpe_tokenizer.encode(sample_src)\nprint(f\"   Tokens: {len(encoded_src['input_ids'])}\")\nprint(f\"   IDs (first 20): {encoded_src['input_ids'][:20]}\")\ndecoded_src = bpe_tokenizer.decode(encoded_src['input_ids'])\nprint(f\"   Decoded: {decoded_src[:150]}...\")\n\nprint(f\"\\n2. Target text:\")\nprint(f\"   {sample_tgt[:150]}...\")\nencoded_tgt = bpe_tokenizer.encode(sample_tgt)\nprint(f\"   Tokens: {len(encoded_tgt['input_ids'])}\")\nprint(f\"   IDs (first 20): {encoded_tgt['input_ids'][:20]}\")\ndecoded_tgt = bpe_tokenizer.decode(encoded_tgt['input_ids'])\nprint(f\"   Decoded: {decoded_tgt[:150]}...\")\n\n# Show some learned merges\nprint(f\"\\n3. Sample learned merges (first 10):\")\nfor i, merge in enumerate(bpe_tokenizer.merges[:10]):\n    print(f\"   {i+1}. '{merge[0]}' + '{merge[1]}' -> '{merge[0] + merge[1]}'\")\n'''\n# ============================================================================\n# LOAD PRETRAINED BPE TOKENIZER\n# ============================================================================\n\nprint(\"\\n\" + \"=\"*80)\nprint(\"LOADING BPE TOKENIZER\")\nprint(\"=\"*80)\n\nTOKENIZER_PATH = \"/kaggle/input/bpe-tock/bpe_tokenizer.pkl\"\n\nbpe_tokenizer = SimpleBPETokenizer.load(TOKENIZER_PATH)\n\nprint(f\"Vocabulary size: {len(bpe_tokenizer.token_to_id)}\")\nprint(f\"Number of merges: {len(bpe_tokenizer.merges)}\")\n\n\n# ============================================================================\n# DATASET FOR PYTORCH\n# ============================================================================\n\nMAX_SRC = 256\nMAX_TGT = 128\n\nclass MTDatasetBPE(Dataset):\n    def __init__(self, df, tokenizer):\n        self.src = df[\"src\"].tolist()\n        self.tgt = df[\"tgt\"].tolist()\n        self.tokenizer = tokenizer\n\n    def __len__(self):\n        return len(self.src)\n\n    def __getitem__(self, idx):\n        # Encode source\n        src_enc = self.tokenizer.encode(\n            self.src[idx],\n            add_special_tokens=True,\n            max_length=MAX_SRC,\n            truncation=True,\n            padding=False\n        )\n        \n        # Encode target\n        tgt_enc = self.tokenizer.encode(\n            self.tgt[idx],\n            add_special_tokens=True,\n            max_length=MAX_TGT,\n            truncation=True,\n            padding=False\n        )\n        \n        return {\n            'input_ids': torch.tensor(src_enc['input_ids'], dtype=torch.long),\n            'attention_mask': torch.tensor(src_enc['attention_mask'], dtype=torch.long),\n            'labels': torch.tensor(tgt_enc['input_ids'], dtype=torch.long)\n        }\n\n\n# Split data\nprint(\"\\n\" + \"=\"*80)\nprint(\"CREATING DATASETS\")\nprint(\"=\"*80)\n\nval_frac = 0.08\npairs_df = pairs_df.sample(frac=1.0, random_state=SEED).reset_index(drop=True)\nn_val = max(1, int(len(pairs_df) * val_frac))\nval_df = pairs_df.iloc[:n_val].reset_index(drop=True)\ntrn_df = pairs_df.iloc[n_val:].reset_index(drop=True)\n\ntrain_ds = MTDatasetBPE(trn_df, bpe_tokenizer)\nval_ds = MTDatasetBPE(val_df, bpe_tokenizer)\n\nprint(f\"Training samples: {len(train_ds)}\")\nprint(f\"Validation samples: {len(val_ds)}\")\n\n# Test dataset\nsample = train_ds[0]\nprint(f\"\\nSample batch shape:\")\nprint(f\"  input_ids: {sample['input_ids'].shape}\")\nprint(f\"  attention_mask: {sample['attention_mask'].shape}\")\nprint(f\"  labels: {sample['labels'].shape}\")\n\nprint(\"\\n\" + \"=\"*80)\nprint(\"BPE TOKENIZER READY FOR TRAINING!\")\nprint(\"=\"*80)\nprint(f\"Vocabulary size: {len(bpe_tokenizer.token_to_id)}\")\nprint(f\"Number of merges: {len(bpe_tokenizer.merges)}\")\nprint(f\"Tokenizer saved to: bpe_tokenizer.pkl\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:51:51.448713Z","iopub.execute_input":"2025-12-22T10:51:51.449250Z","iopub.status.idle":"2025-12-22T10:51:52.213952Z","shell.execute_reply.started":"2025-12-22T10:51:51.449223Z","shell.execute_reply":"2025-12-22T10:51:52.213165Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pickle\n\nwith open(\"/kaggle/input/bpe-tock/bpe_tokenizer.pkl\", \"rb\") as f:\n    tokenizer = pickle.load(f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:51:52.216067Z","iopub.execute_input":"2025-12-22T10:51:52.216402Z","iopub.status.idle":"2025-12-22T10:51:52.227737Z","shell.execute_reply.started":"2025-12-22T10:51:52.216379Z","shell.execute_reply":"2025-12-22T10:51:52.226970Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, re, math, random\nimport numpy as np\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.utils.rnn import pad_sequence\nimport pickle\n\nSEED = 1\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)\n\n# ============================================================================\n# TRANSFORMER SEQ2SEQ MODEL\n# ============================================================================\n\nclass TransformerTranslationModel(nn.Module):\n    \"\"\"\n    Transformer-based Seq2Seq model for translation.\n    \"\"\"\n    \n    def __init__(self, vocab_size, d_model=512, nhead=8, \n                 num_encoder_layers=6, num_decoder_layers=6,\n                 dim_feedforward=2048, dropout=0.1, max_len=512,\n                 pad_token_id=0):\n        super().__init__()\n        \n        self.d_model = d_model\n        self.pad_token_id = pad_token_id\n        \n        # Embeddings\n        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_token_id)\n        self.pos_encoder = PositionalEncoding(d_model, dropout, max_len)\n        \n        # Transformer\n        self.transformer = nn.Transformer(\n            d_model=d_model,\n            nhead=nhead,\n            num_encoder_layers=num_encoder_layers,\n            num_decoder_layers=num_decoder_layers,\n            dim_feedforward=dim_feedforward,\n            dropout=dropout,\n            batch_first=True\n        )\n        \n        # Output projection\n        self.fc_out = nn.Linear(d_model, vocab_size)\n        \n        self._init_weights()\n    \n    def _init_weights(self):\n        \"\"\"Initialize weights.\"\"\"\n        for p in self.parameters():\n            if p.dim() > 1:\n                nn.init.xavier_uniform_(p)\n    \n    def generate_square_subsequent_mask(self, sz):\n        \"\"\"Generate causal mask for decoder.\"\"\"\n        mask = torch.triu(torch.ones(sz, sz), diagonal=1).bool()\n        return mask\n    \n    def forward(self, src, tgt, src_key_padding_mask=None, \n                tgt_key_padding_mask=None):\n        \"\"\"\n        Args:\n            src: [batch, src_len]\n            tgt: [batch, tgt_len]\n            src_key_padding_mask: [batch, src_len] (True for padding)\n            tgt_key_padding_mask: [batch, tgt_len] (True for padding)\n        \"\"\"\n        # Embed and add positional encoding\n        src_emb = self.pos_encoder(self.embedding(src) * math.sqrt(self.d_model))\n        tgt_emb = self.pos_encoder(self.embedding(tgt) * math.sqrt(self.d_model))\n        \n        # Create causal mask for decoder\n        tgt_mask = self.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)\n        \n        # Forward through transformer\n        output = self.transformer(\n            src_emb, \n            tgt_emb,\n            tgt_mask=tgt_mask,\n            src_key_padding_mask=src_key_padding_mask,\n            tgt_key_padding_mask=tgt_key_padding_mask\n        )\n        \n        # Project to vocabulary\n        logits = self.fc_out(output)\n        \n        return logits\n\n\nclass PositionalEncoding(nn.Module):\n    \"\"\"Positional encoding for transformer.\"\"\"\n    \n    def __init__(self, d_model, dropout=0.1, max_len=5000):\n        super().__init__()\n        self.dropout = nn.Dropout(p=dropout)\n        \n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * \n                            (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0)\n        self.register_buffer('pe', pe)\n    \n    def forward(self, x):\n        x = x + self.pe[:, :x.size(1), :]\n        return self.dropout(x)\n\n\n# ============================================================================\n# DATASET AND COLLATOR\n# ============================================================================\n\nclass TranslationDataset(Dataset):\n    \"\"\"Dataset for translation with BPE tokenizer.\"\"\"\n    \n    def __init__(self, df, tokenizer, max_src_len=256, max_tgt_len=128):\n        self.src = df[\"src\"].tolist()\n        self.tgt = df[\"tgt\"].tolist()\n        self.tokenizer = tokenizer\n        self.max_src_len = max_src_len\n        self.max_tgt_len = max_tgt_len\n    \n    def __len__(self):\n        return len(self.src)\n    \n    def __getitem__(self, idx):\n        # Encode source\n        src_enc = self.tokenizer.encode(\n            self.src[idx],\n            add_special_tokens=True,\n            max_length=self.max_src_len,\n            truncation=True\n        )\n        \n        # Encode target\n        tgt_enc = self.tokenizer.encode(\n            self.tgt[idx],\n            add_special_tokens=True,\n            max_length=self.max_tgt_len,\n            truncation=True\n        )\n        \n        return {\n            'src_ids': torch.tensor(src_enc['input_ids'], dtype=torch.long),\n            'src_mask': torch.tensor(src_enc['attention_mask'], dtype=torch.long),\n            'tgt_ids': torch.tensor(tgt_enc['input_ids'], dtype=torch.long),\n            'tgt_mask': torch.tensor(tgt_enc['attention_mask'], dtype=torch.long),\n        }\n\n\ndef collate_fn(batch, pad_token_id=0):\n    \"\"\"Collate function for batching with padding.\"\"\"\n    src_ids = [item['src_ids'] for item in batch]\n    tgt_ids = [item['tgt_ids'] for item in batch]\n    \n    # Pad sequences\n    src_ids_padded = pad_sequence(src_ids, batch_first=True, \n                                   padding_value=pad_token_id)\n    tgt_ids_padded = pad_sequence(tgt_ids, batch_first=True, \n                                   padding_value=pad_token_id)\n    \n    # Create padding masks (True for padding positions)\n    src_key_padding_mask = (src_ids_padded == pad_token_id)\n    tgt_key_padding_mask = (tgt_ids_padded == pad_token_id)\n    \n    return {\n        'src': src_ids_padded,\n        'tgt_input': tgt_ids_padded[:, :-1],  # Remove last token for decoder input\n        'tgt_output': tgt_ids_padded[:, 1:],  # Remove first token for target\n        'src_key_padding_mask': src_key_padding_mask,\n        'tgt_key_padding_mask': tgt_key_padding_mask[:, :-1],\n    }\n\n\n# ============================================================================\n# TRAINING LOOP\n# ============================================================================\n\ndef train_epoch(model, dataloader, optimizer, criterion, device, grad_clip=1.0):\n    \"\"\"Train for one epoch.\"\"\"\n    model.train()\n    total_loss = 0\n    \n    for batch_idx, batch in enumerate(dataloader):\n        # Move to device\n        src = batch['src'].to(device)\n        tgt_input = batch['tgt_input'].to(device)\n        tgt_output = batch['tgt_output'].to(device)\n        src_mask = batch['src_key_padding_mask'].to(device)\n        tgt_mask = batch['tgt_key_padding_mask'].to(device)\n        \n        # Forward pass\n        optimizer.zero_grad()\n        logits = model(src, tgt_input, src_mask, tgt_mask)\n        \n        # Calculate loss (ignore padding)\n        loss = criterion(\n            logits.reshape(-1, logits.size(-1)),\n            tgt_output.reshape(-1)\n        )\n        \n        # Backward pass\n        loss.backward()\n        \n        # Gradient clipping\n        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n        \n        optimizer.step()\n        \n        total_loss += loss.item()\n        \n        if (batch_idx + 1) % 50 == 0:\n            avg_loss = total_loss / (batch_idx + 1)\n            print(f\"  Batch {batch_idx+1}/{len(dataloader)}, Loss: {avg_loss:.4f}\")\n    \n    return total_loss / len(dataloader)\n\n\ndef evaluate(model, dataloader, criterion, device):\n    \"\"\"Evaluate the model.\"\"\"\n    model.eval()\n    total_loss = 0\n    \n    with torch.no_grad():\n        for batch in dataloader:\n            src = batch['src'].to(device)\n            tgt_input = batch['tgt_input'].to(device)\n            tgt_output = batch['tgt_output'].to(device)\n            src_mask = batch['src_key_padding_mask'].to(device)\n            tgt_mask = batch['tgt_key_padding_mask'].to(device)\n            \n            logits = model(src, tgt_input, src_mask, tgt_mask)\n            \n            loss = criterion(\n                logits.reshape(-1, logits.size(-1)),\n                tgt_output.reshape(-1)\n            )\n            \n            total_loss += loss.item()\n    \n    return total_loss / len(dataloader)\n\n\n# ============================================================================\n# INFERENCE / TRANSLATION\n# ============================================================================\n\ndef translate(model, tokenizer, src_text, device, max_len=128, beam_size=4):\n    \"\"\"\n    Translate a source text using beam search.\n    \n    Args:\n        model: Trained translation model\n        tokenizer: BPE tokenizer\n        src_text: Source text string\n        device: torch device\n        max_len: Maximum generation length\n        beam_size: Beam size for beam search\n    \"\"\"\n    model.eval()\n    \n    # Encode source\n    src_enc = tokenizer.encode(src_text, add_special_tokens=True)\n    src_ids = torch.tensor([src_enc['input_ids']]).to(device)\n    src_mask = (src_ids == tokenizer.pad_token_id)\n    \n    # Start with BOS token\n    tgt_ids = torch.tensor([[tokenizer.bos_token_id]]).to(device)\n    \n    with torch.no_grad():\n        for _ in range(max_len):\n            tgt_mask = (tgt_ids == tokenizer.pad_token_id)\n            \n            # Forward pass\n            logits = model(src_ids, tgt_ids, src_mask, tgt_mask)\n            \n            # Get next token (greedy for simplicity)\n            next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)\n            \n            # Append to sequence\n            tgt_ids = torch.cat([tgt_ids, next_token], dim=1)\n            \n            # Stop if EOS token\n            if next_token.item() == tokenizer.eos_token_id:\n                break\n    \n    # Decode\n    translation = tokenizer.decode(tgt_ids[0].cpu().tolist(), \n                                   skip_special_tokens=True)\n    return translation\n\n\ndef translate_long_text(model, tokenizer, src_text, device, \n                        chunk_size=200, max_len=128):\n    \"\"\"\n    Translate long text by chunking.\n    \n    Args:\n        model: Trained translation model\n        tokenizer: BPE tokenizer\n        src_text: Long source text\n        device: torch device\n        chunk_size: Characters per chunk\n        max_len: Maximum generation length per chunk\n    \"\"\"\n    # Split into chunks\n    chunks = []\n    for i in range(0, len(src_text), chunk_size):\n        chunk = src_text[i:i+chunk_size]\n        if chunk.strip():\n            chunks.append(chunk)\n    \n    # Translate each chunk\n    translations = []\n    for chunk in chunks:\n        translation = translate(model, tokenizer, chunk, device, max_len)\n        translations.append(translation)\n    \n    # Join translations\n    full_translation = \" \".join(translations)\n    \n    # Clean up punctuation\n    if not full_translation.endswith('.'):\n        full_translation += '.'\n    \n    return full_translation\n\n\n# ============================================================================\n# MAIN TRAINING SCRIPT\n# ============================================================================\n\n# ============================================================\n# Load data\n# ============================================================\nprint(\"Loading data...\")\nTRAIN_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\nTEST_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n\ntrain_df = pd.read_csv(TRAIN_PATH)\ntest_df = pd.read_csv(TEST_PATH)\n\n# ============================================================\n# Load BPE tokenizer\n# ============================================================\nprint(\"\\nLoading BPE tokenizer...\")\n\nclass SimpleBPETokenizer:\n    pass\n\ntry:\n    with open(\"/kaggle/input/bpe-tock/bpe_tokenizer.pkl\", \"rb\") as f:\n        tokenizer_data = pickle.load(f)\n\n    tokenizer = SimpleBPETokenizer()\n    tokenizer.token_to_id = tokenizer_data[\"token_to_id\"]\n    tokenizer.id_to_token = tokenizer_data[\"id_to_token\"]\n    tokenizer.merges = tokenizer_data[\"merges\"]\n    tokenizer.vocab_size = tokenizer_data[\"vocab_size\"]\n\n    for key, (token, token_id) in tokenizer_data[\"special_tokens\"].items():\n        setattr(tokenizer, f\"{key}_token\", token)\n        setattr(tokenizer, f\"{key}_token_id\", token_id)\n\n    def encode(self, text, add_special_tokens=True, max_length=None, truncation=True):\n        if pd.isna(text):\n            text = \"\"\n        text = str(text).strip()\n\n        def tokenize_word(word):\n            if not word:\n                return []\n            tokens = list(word) + [\"</w>\"]\n            for a, b in self.merges:\n                i = 0\n                new_tokens = []\n                while i < len(tokens):\n                    if i < len(tokens) - 1 and tokens[i] == a and tokens[i + 1] == b:\n                        new_tokens.append(a + b)\n                        i += 2\n                    else:\n                        new_tokens.append(tokens[i])\n                        i += 1\n                tokens = new_tokens\n                if len(tokens) == 1:\n                    break\n            return tokens\n\n        all_tokens = []\n        for word in text.split():\n            all_tokens.extend(tokenize_word(word))\n\n        token_ids = []\n        if add_special_tokens:\n            token_ids.append(self.bos_token_id)\n\n        for tok in all_tokens:\n            token_ids.append(self.token_to_id.get(tok, self.unk_token_id))\n\n        if add_special_tokens:\n            token_ids.append(self.eos_token_id)\n\n        if max_length and truncation and len(token_ids) > max_length:\n            token_ids = token_ids[: max_length - 1] + [self.eos_token_id]\n\n        attention_mask = [1 if t != self.pad_token_id else 0 for t in token_ids]\n\n        return {\"input_ids\": token_ids, \"attention_mask\": attention_mask}\n\n    def decode(self, token_ids, skip_special_tokens=True):\n        if isinstance(token_ids, torch.Tensor):\n            token_ids = token_ids.tolist()\n\n        tokens = []\n        for tid in token_ids:\n            if skip_special_tokens and tid in {\n                self.pad_token_id,\n                self.bos_token_id,\n                self.eos_token_id,\n                self.unk_token_id,\n            }:\n                continue\n            tokens.append(self.id_to_token.get(tid, self.unk_token))\n\n        text = \"\".join(tokens).replace(\"</w>\", \" \")\n        return text.strip()\n\n    tokenizer.encode = lambda text, **kw: encode(tokenizer, text, **kw)\n    tokenizer.decode = lambda ids, **kw: decode(tokenizer, ids, **kw)\n\n    print(\"✓ Tokenizer loaded\")\n    print(f\"Vocabulary size: {len(tokenizer.token_to_id)}\")\n\nexcept Exception as e:\n    print(f\"Tokenizer load failed: {e}\")\n\n# ============================================================\n# Build aligned pairs (only for sample inference sanity check)\n# ============================================================\nprint(\"\\nPreparing aligned pairs (for sample check)...\")\n\ndef clean_ws(x):\n    return \"\" if pd.isna(x) else re.sub(r\"\\s+\", \" \", str(x)).strip()\n\ndef split_src_transliteration(doc):\n    doc = \"\" if pd.isna(doc) else str(doc)\n    parts = [clean_ws(p) for p in re.split(r\"\\n+\", doc) if clean_ws(p)]\n    return parts if parts else [clean_ws(doc)]\n\ndef split_tgt_english(doc):\n    return [clean_ws(s) for s in re.split(r\"(?<=[.!?])\\s+\", clean_ws(doc)) if clean_ws(s)]\n\ndef tok_len(s):\n    return max(1, len(s.split()))\n\ndef align_monotone(src, tgt):\n    pairs = []\n    i = j = 0\n    while i < len(src) and j < len(tgt):\n        pairs.append((src[i], tgt[j]))\n        i += 1\n        j += 1\n    return pairs\n\npairs = []\nfor _, row in train_df.iterrows():\n    s = split_src_transliteration(row[\"transliteration\"])\n    t = split_tgt_english(row[\"translation\"])\n    if s and t:\n        pairs.extend(align_monotone(s, t))\n\npairs_df = pd.DataFrame(pairs, columns=[\"src\", \"tgt\"]).drop_duplicates()\nprint(f\"Total aligned pairs: {len(pairs_df)}\")\n\n# ============================================================\n# Initialize model and load BEST checkpoint\n# ============================================================\nprint(\"\\nLoading trained model...\")\n\nvocab_size = len(tokenizer.token_to_id)\n\nbpe_model = TransformerTranslationModel(\n    vocab_size=vocab_size,\n    d_model=256,\n    nhead=8,\n    num_encoder_layers=4,\n    num_decoder_layers=4,\n    dim_feedforward=1024,\n    dropout=0.1,\n    pad_token_id=tokenizer.pad_token_id,\n).to(DEVICE)\n\nbpe_model.load_state_dict(\n    torch.load(\"/kaggle/input/best-translation-model/best_translation_model.pt\", map_location=DEVICE)\n)\nbpe_model.eval()\n\nprint(\"✓ Best model loaded\")\nprint(f\"Model parameters: {sum(p.numel() for p in bpe_model.parameters()):,}\")\n\n# ============================================================\n# Sample inference check\n# ============================================================\nprint(\"\\n\" + \"=\" * 80)\nprint(\"SAMPLE TRANSLATION\")\nprint(\"=\" * 80)\n\nsample_src = pairs_df[\"src\"].iloc[0]\nsample_tgt = pairs_df[\"tgt\"].iloc[0]\n\nwith torch.no_grad():\n    pred = translate(bpe_model, tokenizer, sample_src, DEVICE)\n\nprint(f\"Source: {sample_src}\")\nprint(f\"Target: {sample_tgt}\")\nprint(f\"Prediction: {pred}\")\n\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:51:52.228924Z","iopub.execute_input":"2025-12-22T10:51:52.229204Z","iopub.status.idle":"2025-12-22T10:52:00.651025Z","shell.execute_reply.started":"2025-12-22T10:51:52.229181Z","shell.execute_reply":"2025-12-22T10:52:00.650157Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install /kaggle/input/whl2eval/evaluate-0.4.6-py3-none-any.whl\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:52:00.652171Z","iopub.execute_input":"2025-12-22T10:52:00.652446Z","iopub.status.idle":"2025-12-22T10:52:04.137730Z","shell.execute_reply.started":"2025-12-22T10:52:00.652410Z","shell.execute_reply":"2025-12-22T10:52:04.136980Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/morewhlfiles/portalocker-3.2.0-py3-none-any.whl\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:52:04.139134Z","iopub.execute_input":"2025-12-22T10:52:04.139459Z","iopub.status.idle":"2025-12-22T10:52:07.308166Z","shell.execute_reply.started":"2025-12-22T10:52:04.139429Z","shell.execute_reply":"2025-12-22T10:52:07.307432Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/whl2eval/sacrebleu-2.5.1-py3-none-any.whl\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T10:52:07.309502Z","iopub.execute_input":"2025-12-22T10:52:07.310175Z","iopub.status.idle":"2025-12-22T10:52:10.605027Z","shell.execute_reply.started":"2025-12-22T10:52:07.310144Z","shell.execute_reply":"2025-12-22T10:52:10.604161Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.metrics.pairwise import cosine_similarity\nimport numpy as np\nimport re\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom tqdm import tqdm\ntry:\n    from transformers import get_linear_schedule_with_warmup\nexcept:\n    print(\"Warning: transformers not available, using simple scheduler\")\n    get_linear_schedule_with_warmup = None\n\n# Set device - use cuda:0 explicitly to avoid multi-GPU issues\nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {DEVICE}\")\n\n# Import sacrebleu for metrics if available\ntry:\n    import sacrebleu\n    SACREBLEU_AVAILABLE = True\nexcept:\n    print(\"Warning: sacrebleu not available, metrics will be approximate\")\n    SACREBLEU_AVAILABLE = False\n\ndef compute_bleu_simple(predictions, references):\n    \"\"\"Simple BLEU approximation using word overlap\"\"\"\n    from collections import Counter\n    \n    total_score = 0\n    for pred, ref_list in zip(predictions, references):\n        ref = ref_list[0] if isinstance(ref_list, list) else ref_list\n        \n        pred_words = pred.lower().split()\n        ref_words = ref.lower().split()\n        \n        if len(pred_words) == 0:\n            total_score += 0\n            continue\n        \n        # Simple precision calculation\n        pred_counter = Counter(pred_words)\n        ref_counter = Counter(ref_words)\n        \n        overlap = sum((pred_counter & ref_counter).values())\n        precision = overlap / len(pred_words) if len(pred_words) > 0 else 0\n        \n        # Brevity penalty\n        bp = min(1.0, np.exp(1 - len(ref_words) / max(len(pred_words), 1)))\n        \n        total_score += bp * precision\n    \n    return total_score / len(predictions)\n\ndef compute_chrf_simple(predictions, references):\n    \"\"\"Simple character F-score approximation\"\"\"\n    total_score = 0\n    \n    for pred, ref_list in zip(predictions, references):\n        ref = ref_list[0] if isinstance(ref_list, list) else ref_list\n        \n        pred_chars = set(pred.lower())\n        ref_chars = set(ref.lower())\n        \n        if len(pred_chars) == 0 or len(ref_chars) == 0:\n            total_score += 0\n            continue\n        \n        overlap = len(pred_chars & ref_chars)\n        precision = overlap / len(pred_chars)\n        recall = overlap / len(ref_chars)\n        \n        f_score = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0\n        total_score += f_score\n    \n    return total_score / len(predictions)\n\ndef compute_geometric_mean_score(predictions, references):\n    \"\"\"Compute geometric mean of BLEU and chrF++ scores\"\"\"\n    if SACREBLEU_AVAILABLE:\n        # Use sacrebleu for accurate metrics\n        refs_transposed = [[ref[0] if isinstance(ref, list) else ref for ref in references]]\n        bleu = sacrebleu.corpus_bleu(predictions, refs_transposed)\n        chrf = sacrebleu.corpus_chrf(predictions, refs_transposed)\n        \n        bleu_score = bleu.score / 100.0\n        chrf_score = chrf.score / 100.0\n    else:\n        # Use simple approximations\n        bleu_score = compute_bleu_simple(predictions, references)\n        chrf_score = compute_chrf_simple(predictions, references)\n    \n    geometric_mean = np.sqrt(bleu_score * chrf_score)\n    return geometric_mean, bleu_score, chrf_score\n\ndef make_one_sentence(en: str) -> str:\n    en = clean_ws(en)\n    if not en:\n        return \"...\"\n    parts = re.split(r\"(?<=[.!?])\\s+\", en)\n    parts = [p.strip() for p in parts if p.strip()]\n    # remove ending punctuation to avoid multiple sentences\n    parts = [re.sub(r\"[.!?]+\\s*$\", \"\", p).strip() for p in parts]\n    parts = [p for p in parts if p]\n    out = \"; \".join(parts).strip()\n    if not out.endswith(\".\"):\n        out += \".\"\n    return out\n\n# Clean training sources\ntrain_src = [clean_ws(t) for t in train_df[\"transliteration\"].tolist()]\ntrain_tgt = [clean_ws(t) for t in train_df[\"translation\"].tolist()]\n\n# ============================================================================\n# STEP 1: Fine-tune the simple BPE model on training data\n# ============================================================================\nprint(\"=\" * 80)\nprint(\"STEP 1: Fine-tuning simple BPE model\")\nprint(\"=\" * 80)\n\nclass TranslationDataset(Dataset):\n    def __init__(self, sources, targets, tokenizer, max_length=128):\n        self.sources = sources\n        self.targets = targets\n        self.tokenizer = tokenizer\n        self.max_length = max_length\n        \n        # Determine pad token id\n        if hasattr(tokenizer, 'pad_token_id') and tokenizer.pad_token_id is not None:\n            self.pad_id = tokenizer.pad_token_id\n        elif hasattr(tokenizer, 'vocab') and '<pad>' in tokenizer.vocab:\n            self.pad_id = tokenizer.vocab['<pad>']\n        else:\n            # Use 0 as default pad id\n            self.pad_id = 0\n    \n    def __len__(self):\n        return len(self.sources)\n    \n    def __getitem__(self, idx):\n        source = self.sources[idx]\n        target = self.targets[idx]\n        \n        # Tokenize using encode method if available, otherwise use text_to_sequence\n        if hasattr(self.tokenizer, 'encode'):\n            src_result = self.tokenizer.encode(source)\n            tgt_result = self.tokenizer.encode(target)\n        elif hasattr(self.tokenizer, 'text_to_sequence'):\n            src_result = self.tokenizer.text_to_sequence(source)\n            tgt_result = self.tokenizer.text_to_sequence(target)\n        else:\n            raise AttributeError(\"Tokenizer must have 'encode' or 'text_to_sequence' method\")\n        \n        # Extract token IDs from result (handle dict, tensor, or list)\n        if isinstance(src_result, dict):\n            src_tokens = src_result['input_ids'] if 'input_ids' in src_result else src_result['ids']\n        else:\n            src_tokens = src_result\n            \n        if isinstance(tgt_result, dict):\n            tgt_tokens = tgt_result['input_ids'] if 'input_ids' in tgt_result else tgt_result['ids']\n        else:\n            tgt_tokens = tgt_result\n        \n        # Convert to list if needed\n        if isinstance(src_tokens, torch.Tensor):\n            src_tokens = src_tokens.tolist()\n        if isinstance(tgt_tokens, torch.Tensor):\n            tgt_tokens = tgt_tokens.tolist()\n        \n        # Ensure we have a list\n        if not isinstance(src_tokens, list):\n            src_tokens = list(src_tokens)\n        if not isinstance(tgt_tokens, list):\n            tgt_tokens = list(tgt_tokens)\n        \n        # Truncate if needed\n        if len(src_tokens) > self.max_length:\n            src_tokens = src_tokens[:self.max_length]\n        if len(tgt_tokens) > self.max_length:\n            tgt_tokens = tgt_tokens[:self.max_length]\n        \n        # Pad sequences\n        src_len = len(src_tokens)\n        tgt_len = len(tgt_tokens)\n        \n        src_padded = src_tokens + [self.pad_id] * (self.max_length - src_len)\n        tgt_padded = tgt_tokens + [self.pad_id] * (self.max_length - tgt_len)\n        \n        # Create attention mask (1 for real tokens, 0 for padding)\n        src_mask = [1] * src_len + [0] * (self.max_length - src_len)\n        \n        return {\n            'input_ids': torch.tensor(src_padded, dtype=torch.long),\n            'attention_mask': torch.tensor(src_mask, dtype=torch.long),\n            'labels': torch.tensor(tgt_padded, dtype=torch.long)\n        }\n\ndef finetune_bpe_model(model, tokenizer, train_sources, train_targets, epochs=3, batch_size=8, lr=5e-5):\n    \"\"\"Fine-tune the BPE model on training data\"\"\"\n    model = model.to(DEVICE)\n    model.train()\n    \n    # Create dataset and dataloader\n    dataset = TranslationDataset(train_sources, train_targets, tokenizer)\n    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n    \n    # Optimizer and scheduler\n    optimizer = AdamW(model.parameters(), lr=lr)\n    total_steps = len(dataloader) * epochs\n    \n    if get_linear_schedule_with_warmup is not None:\n        scheduler = get_linear_schedule_with_warmup(\n            optimizer,\n            num_warmup_steps=int(0.1 * total_steps),\n            num_training_steps=total_steps\n        )\n    else:\n        # Simple step scheduler fallback\n        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=len(dataloader), gamma=0.9)\n    \n    # Define loss function\n    criterion = torch.nn.CrossEntropyLoss(ignore_index=dataset.pad_id)\n    \n    print(f\"Training for {epochs} epochs with {len(dataloader)} batches per epoch\")\n    \n    for epoch in range(epochs):\n        total_loss = 0\n        progress_bar = tqdm(dataloader, desc=f\"Epoch {epoch+1}/{epochs}\")\n        \n        for batch in progress_bar:\n            input_ids = batch['input_ids'].to(DEVICE)\n            attention_mask = batch['attention_mask'].to(DEVICE)\n            labels = batch['labels'].to(DEVICE)\n            \n            # Forward pass - adjust based on model's actual forward signature\n            # Most transformer models expect (src, tgt) or similar\n            try:\n                # Try standard seq2seq format\n                outputs = model(input_ids, labels[:, :-1])  # Teacher forcing\n                \n                # Get logits and compute loss\n                if isinstance(outputs, dict) and 'logits' in outputs:\n                    logits = outputs['logits']\n                elif isinstance(outputs, tuple):\n                    logits = outputs[0]\n                else:\n                    logits = outputs\n                \n                # Reshape for loss calculation\n                loss = criterion(\n                    logits.reshape(-1, logits.size(-1)),\n                    labels[:, 1:].reshape(-1)\n                )\n            except Exception as e:\n                # Fallback: try direct forward pass\n                try:\n                    outputs = model(input_ids)\n                    if isinstance(outputs, dict) and 'logits' in outputs:\n                        logits = outputs['logits']\n                    elif isinstance(outputs, tuple):\n                        logits = outputs[0]\n                    else:\n                        logits = outputs\n                    \n                    loss = criterion(\n                        logits.reshape(-1, logits.size(-1)),\n                        labels.reshape(-1)\n                    )\n                except Exception as e2:\n                    print(f\"Error in forward pass: {e2}\")\n                    print(\"Skipping batch due to model signature mismatch\")\n                    continue\n            \n            total_loss += loss.item()\n            \n            # Backward pass\n            optimizer.zero_grad()\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            optimizer.step()\n            scheduler.step()\n            \n            progress_bar.set_postfix({'loss': loss.item()})\n        \n        avg_loss = total_loss / len(dataloader)\n        print(f\"Epoch {epoch+1} average loss: {avg_loss:.4f}\")\n        \n        # Clear cache\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n    \n    model.eval()\n    return model\n\n# Fine-tune BPE model\nbpe_model_finetuned = finetune_bpe_model(\n    bpe_model, \n    tokenizer, \n    train_src, \n    train_tgt,\n    epochs=3,\n    batch_size=8\n)\n\n# ============================================================================\n# STEP 2: Fine-tune the LLM with reasoning on training data\n# ============================================================================\nprint(\"=\" * 80)\nprint(\"STEP 2: Fine-tuning LLM with reasoning\")\nprint(\"=\" * 80)\n\n# Setup TF-IDF for retrieval\nvectorizer = TfidfVectorizer(\n    analyzer=\"char\",\n    ngram_range=(3, 6),\n    min_df=2,\n)\n\ntrain_tfidf = vectorizer.fit_transform(train_src)\n\ndef get_top_k_examples(query: str, k: int = 1, exclude_idx=None):\n    query = clean_ws(query)\n    if not query:\n        return []\n\n    q_vec = vectorizer.transform([query])\n    sims = cosine_similarity(q_vec, train_tfidf)[0]\n    \n    # Exclude the query itself if it's in training set\n    if exclude_idx is not None:\n        sims[exclude_idx] = -1\n\n    top_idx = np.argsort(sims)[-k:][::-1]\n\n    examples = []\n    for i in top_idx:\n        examples.append((train_src[i], train_tgt[i]))\n\n    return examples\n\ndef build_reasoning_prompt(src_text: str, examples, simple_model_translation: str) -> str:\n    parts = [\n        \"You are an expert translator of ancient transliterations into English.\",\n        \"Below are similar examples with their translations.\",\n        \"Study the patterns and vocabulary, then translate the new text.\",\n        \"\"\n    ]\n\n    for i, (source, translation) in enumerate(examples, start=1):\n        parts.extend([\n            f\"Example {i}:\",\n            f\"Source: {source}\",\n            f\"Translation: {translation}\",\n            \"\"\n        ])\n\n    parts.extend([\n        \"Also this is the translation of a simpler model.\",\n        \"\",\n        \"Source:\",\n        simple_model_translation,\n        \"\",\n        \"Now translate the following text into clear English.\",\n        \"\",\n        \"Source:\",\n        src_text,\n        \"\",\n        \"Provide your translation between <start_of_translation> and <end_of_translation> tags.\",\n        \"\",\n        \"Translation:\"\n    ])\n\n    return \"\\n\".join(parts)\n\nclass LLMReasoningDataset(Dataset):\n    def __init__(self, sources, targets, bpe_model, tokenizer, device):\n        self.sources = sources\n        self.targets = targets\n        self.bpe_model = bpe_model\n        self.tokenizer = tokenizer\n        self.device = device\n        \n        # Pre-generate BPE translations\n        print(\"Generating BPE translations for training...\")\n        self.bpe_translations = []\n        for src in tqdm(sources):\n            bpe_trans = translate_long_text(bpe_model, tokenizer, src, device)\n            self.bpe_translations.append(bpe_trans)\n    \n    def __len__(self):\n        return len(self.sources)\n    \n    def __getitem__(self, idx):\n        source = self.sources[idx]\n        target = self.targets[idx]\n        bpe_trans = self.bpe_translations[idx]\n        \n        # Get similar examples (excluding current one)\n        examples = get_top_k_examples(source, k=1, exclude_idx=idx)\n        \n        # Build prompt\n        prompt = build_reasoning_prompt(source, examples, bpe_trans)\n        \n        # Build target with tags\n        full_target = f\"<start_of_translation>{target}<end_of_translation>\"\n        \n        return prompt, full_target\n\ndef finetune_llm_with_reasoning(model, tokenizer, train_sources, train_targets, \n                                 bpe_model, bpe_tokenizer, epochs=2, batch_size=4, lr=2e-5):\n    \"\"\"Fine-tune LLM with reasoning prompts\"\"\"\n    model.train()\n    \n    # Create dataset\n    dataset = LLMReasoningDataset(train_sources, train_targets, bpe_model, bpe_tokenizer, DEVICE)\n    \n    # Since LLM training requires custom handling, we'll do manual batching\n    optimizer = AdamW(model.parameters(), lr=lr)\n    \n    print(f\"Training LLM for {epochs} epochs with {len(dataset)} examples\")\n    \n    for epoch in range(epochs):\n        total_loss = 0\n        num_batches = 0\n        \n        # Shuffle indices\n        indices = torch.randperm(len(dataset))\n        \n        progress_bar = tqdm(range(0, len(dataset), batch_size), desc=f\"Epoch {epoch+1}/{epochs}\")\n        \n        for i in progress_bar:\n            batch_indices = indices[i:i+batch_size]\n            batch_prompts = []\n            batch_targets = []\n            \n            for idx in batch_indices:\n                prompt, target = dataset[idx.item()]\n                batch_prompts.append(prompt)\n                batch_targets.append(target)\n            \n            # Train on batch (implementation depends on your LLM API)\n            # This is a placeholder - adjust based on your model's training API\n            try:\n                # Assuming your model has a train method\n                loss = model.train_batch(batch_prompts, batch_targets)\n                total_loss += loss\n                num_batches += 1\n                \n                progress_bar.set_postfix({'loss': loss})\n            except Exception as e:\n                print(f\"Warning: Could not train batch - {e}\")\n                continue\n            \n            # Clear cache\n            if torch.cuda.is_available() and i % 10 == 0:\n                torch.cuda.empty_cache()\n        \n        avg_loss = total_loss / max(num_batches, 1)\n        print(f\"Epoch {epoch+1} average loss: {avg_loss:.4f}\")\n    \n    return model\n\n# Note: LLM fine-tuning depends on your specific model's API\n# If it doesn't support training, we skip this step\ntry:\n    if hasattr(model, 'train_batch') or hasattr(model, 'train'):\n        model_finetuned = finetune_llm_with_reasoning(\n            model,\n            GLOBAL_TOKENIZER,\n            train_src,\n            train_tgt,\n            bpe_model_finetuned,\n            tokenizer,\n            epochs=2,\n            batch_size=4\n        )\n    else:\n        print(\"LLM does not support fine-tuning, using pre-trained model\")\n        model_finetuned = model\nexcept Exception as e:\n    print(f\"Could not fine-tune LLM: {e}\")\n    model_finetuned = model\n\n# ============================================================================\n# STEP 3: Inference with fine-tuned models\n# ============================================================================\nprint(\"=\" * 80)\nprint(\"STEP 3: Running inference\")\nprint(\"=\" * 80)\n\ndef extract_translation(text: str) -> str:\n    \"\"\"Extract translation content between tags\"\"\"\n    pattern = r'<start_of_translation>(.*?)<end_of_translation>'\n    match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)\n    \n    if match:\n        return match.group(1).strip()\n    else:\n        return text.strip()\n\nMAX_TGT = 1024\n\ndef generate_translation_with_reasoning(src_text: str, row) -> str:\n    src_text = clean_ws(src_text)\n    if not src_text:\n        return \"...\"\n\n    examples = get_top_k_examples(src_text, k=1)\n    bpe_translation = translate_long_text(bpe_model_finetuned, tokenizer, row[\"transliteration\"], DEVICE)\n    prompt = build_reasoning_prompt(src_text, examples, bpe_translation)\n    \n    # Tokenize the prompt\n    prompt_tokens = GLOBAL_TOKENIZER.tokenize(prompt)\n    prompt_length = len(prompt_tokens)\n\n    # Generate output\n    full_output = model_finetuned.generate(\n        prompt,\n        max_length=prompt_length + MAX_TGT,\n    )\n    \n    # Convert tensor to string if needed\n    if hasattr(full_output, 'cpu'):\n        full_output = full_output.cpu()\n    if hasattr(full_output, 'numpy'):\n        full_output = full_output.numpy()\n    \n    # Decode bytes to string if needed\n    if isinstance(full_output, bytes):\n        full_output = full_output.decode('utf-8', errors='ignore')\n    elif isinstance(full_output, np.ndarray):\n        full_output = full_output.item().decode('utf-8', errors='ignore') if isinstance(full_output.item(), bytes) else str(full_output.item())\n    \n    # Extract only generated tokens\n    full_tokens = GLOBAL_TOKENIZER.tokenize(full_output)\n    generated_tokens = full_tokens[prompt_length:]\n    generated_only = GLOBAL_TOKENIZER.detokenize(generated_tokens)\n    \n    # Decode if needed\n    if isinstance(generated_only, bytes):\n        generated_only = generated_only.decode('utf-8', errors='ignore')\n    elif hasattr(generated_only, 'numpy'):\n        generated_only = generated_only.numpy()\n        if isinstance(generated_only, bytes):\n            generated_only = generated_only.decode('utf-8', errors='ignore')\n    \n    print(\"LLM Answer:\", generated_only)\n    \n    # Extract translation\n    extracted = extract_translation(generated_only)\n    \n    return make_one_sentence(extracted)\n\ndef process_batch(texts, batch_size=8):\n    results = []\n    for i in range(0, len(texts), batch_size):\n        batch = texts[i:i+batch_size]\n        batch_results = [\n            generate_translation_with_reasoning(t, row)\n            for t, row in zip(batch, test_df.iloc[i:i+batch_size].to_dict('records'))\n        ]\n        results.extend(batch_results)\n        \n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n    \n    return results\n\n# Process translations\npreds = process_batch(test_df[\"transliteration\"].tolist())\n\n# Evaluate on validation set if available\nif 'translation' in test_df.columns:\n    references = [[ref] for ref in test_df[\"translation\"].tolist()]\n    geometric_mean, bleu, chrf = compute_geometric_mean_score(preds, references)\n    print(f\"\\nEvaluation Metrics:\")\n    print(f\"BLEU Score: {bleu:.4f}\")\n    print(f\"chrF++ Score: {chrf:.4f}\")\n    print(f\"Geometric Mean: {geometric_mean:.4f}\")\n\n# Save submission\nsub = pd.DataFrame({\n    \"id\": test_df[\"id\"],\n    \"translation\": preds\n})\n\nsub.to_csv(\"submission.csv\", index=False)\nprint(\"\\nWrote submission.csv\")\nsub.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:08:16.243668Z","iopub.execute_input":"2025-12-22T11:08:16.244330Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}