{"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":14239027,"sourceType":"datasetVersion","datasetId":9084458},{"sourceId":14243325,"sourceType":"datasetVersion","datasetId":9087187}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-21T09:48:08.078162Z","iopub.execute_input":"2025-12-21T09:48:08.078754Z","iopub.status.idle":"2025-12-21T09:48:13.772683Z","shell.execute_reply.started":"2025-12-21T09:48:08.078717Z","shell.execute_reply":"2025-12-21T09:48:13.771859Z"}},"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)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T09:48:24.307995Z","iopub.execute_input":"2025-12-21T09:48:24.308517Z","iopub.status.idle":"2025-12-21T09:48:24.318773Z","shell.execute_reply.started":"2025-12-21T09:48:24.308489Z","shell.execute_reply":"2025-12-21T09:48:24.317968Z"}},"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\ndef main():\n    # ============================================================\n    # Load data\n    # ============================================================\n    print(\"Loading data...\")\n    TRAIN_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\n    TEST_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n\n    train_df = pd.read_csv(TRAIN_PATH)\n    test_df = pd.read_csv(TEST_PATH)\n\n    # ============================================================\n    # Load BPE tokenizer\n    # ============================================================\n    print(\"\\nLoading BPE tokenizer...\")\n\n    class SimpleBPETokenizer:\n        pass\n\n    try:\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\n    except Exception as e:\n        print(f\"Tokenizer load failed: {e}\")\n        return\n\n    # ============================================================\n    # Build aligned pairs (only for sample inference sanity check)\n    # ============================================================\n    print(\"\\nPreparing aligned pairs (for sample check)...\")\n\n    def clean_ws(x):\n        return \"\" if pd.isna(x) else re.sub(r\"\\s+\", \" \", str(x)).strip()\n\n    def 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\n    def split_tgt_english(doc):\n        return [clean_ws(s) for s in re.split(r\"(?<=[.!?])\\s+\", clean_ws(doc)) if clean_ws(s)]\n\n    def tok_len(s):\n        return max(1, len(s.split()))\n\n    def 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\n    pairs = []\n    for _, 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\n    pairs_df = pd.DataFrame(pairs, columns=[\"src\", \"tgt\"]).drop_duplicates()\n    print(f\"Total aligned pairs: {len(pairs_df)}\")\n\n    # ============================================================\n    # Initialize model and load BEST checkpoint\n    # ============================================================\n    print(\"\\nLoading trained model...\")\n\n    vocab_size = len(tokenizer.token_to_id)\n\n    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\n    model.load_state_dict(\n        torch.load(\"/kaggle/input/best-translation-model/best_translation_model.pt\", map_location=DEVICE)\n    )\n    model.eval()\n\n    print(\"✓ Best model loaded\")\n    print(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\n    # ============================================================\n    # Sample inference check\n    # ============================================================\n    print(\"\\n\" + \"=\" * 80)\n    print(\"SAMPLE TRANSLATION\")\n    print(\"=\" * 80)\n\n    sample_src = pairs_df[\"src\"].iloc[0]\n    sample_tgt = pairs_df[\"tgt\"].iloc[0]\n\n    with torch.no_grad():\n        pred = translate(model, tokenizer, sample_src, DEVICE)\n\n    print(f\"Source: {sample_src}\")\n    print(f\"Target: {sample_tgt}\")\n    print(f\"Prediction: {pred}\")\n\n    # ============================================================\n    # Generate test predictions\n    # ============================================================\n    print(\"\\n\" + \"=\" * 80)\n    print(\"GENERATING TEST PREDICTIONS\")\n    print(\"=\" * 80)\n\n    predictions = []\n\n    with torch.no_grad():\n        for idx, row in test_df.iterrows():\n            translation = translate_long_text(\n                model, tokenizer, row[\"transliteration\"], DEVICE\n            )\n            predictions.append(translation)\n\n            if (idx + 1) % 10 == 0:\n                print(f\"Translated {idx + 1}/{len(test_df)}\")\n\n    submission = pd.DataFrame(\n        {\"id\": test_df[\"id\"], \"translation\": predictions}\n    )\n    submission.to_csv(\"submission.csv\", index=False)\n\n    print(\"\\n✓ Saved submission.csv\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T09:54:02.956254Z","iopub.execute_input":"2025-12-21T09:54:02.956843Z","iopub.status.idle":"2025-12-21T09:54:10.675456Z","shell.execute_reply.started":"2025-12-21T09:54:02.956813Z","shell.execute_reply":"2025-12-21T09:54:10.674864Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}