{"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":15061024,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport math\nimport pandas as pd\nimport re\nimport numpy as np\nfrom torch.utils.data import Dataset, DataLoader\n\n# ============================================================\n# 1. Gated Model Components (Based on arXiv:2505.06708v1)\n# ============================================================\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super().__init__()\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() * (-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        self.register_buffer('pe', pe.unsqueeze(0))\n\n    def forward(self, x):\n        return x + self.pe[:, :x.size(1)]\n\nclass GatedMultiHeadAttention(nn.Module):\n    \"\"\"\n    Implements SDPA Elementwise Gating (G1) as described in the paper.\n    The gating mechanism is formalized as: Y' = Y * sigmoid(X * W_theta)[cite: 53, 105].\n    \"\"\"\n    def __init__(self, d_model, num_heads, dropout=0.1):\n        super().__init__()\n        assert d_model % num_heads == 0\n        self.d_k = d_model // num_heads\n        self.h = num_heads\n        \n        self.q_linear = nn.Linear(d_model, d_model)\n        self.k_linear = nn.Linear(d_model, d_model)\n        self.v_linear = nn.Linear(d_model, d_model)\n        \n        # Learnable gating parameters W_theta [cite: 53]\n        # Maps input X to a gating score for every dimension of every head (Elementwise) [cite: 58, 66]\n        self.gate_linear = nn.Linear(d_model, d_model) \n        \n        self.out_linear = nn.Linear(d_model, d_model)\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, q, k, v, mask=None):\n        bs = q.size(0)\n        \n        # 1. Compute Gating Scores from input X (pre-norm hidden states) [cite: 62]\n        # Paper uses sigmoid activation for the gate [cite: 62, 90]\n        gate = torch.sigmoid(self.gate_linear(q)) # (bs, n, d_model)\n        gate = gate.view(bs, -1, self.h, self.d_k).transpose(1, 2) # (bs, h, n, d_k)\n\n        # 2. Standard Projections\n        q = self.q_linear(q).view(bs, -1, self.h, self.d_k).transpose(1, 2)\n        k = self.k_linear(k).view(bs, -1, self.h, self.d_k).transpose(1, 2)\n        v = self.v_linear(v).view(bs, -1, self.h, self.d_k).transpose(1, 2)\n\n        # 3. Scaled Dot-Product Attention (SDPA)\n        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n        if mask is not None:\n            # Masking for causal attention or padding\n            scores = scores.masked_fill(mask == 0, -1e9)\n        \n        attn = torch.softmax(scores, dim=-1)\n        attn = self.dropout(attn)\n        context = torch.matmul(attn, v) # (bs, h, n, d_k)\n        \n        # 4. Apply Multiplicative Gating [cite: 53, 88]\n        # This modulates the SDPA output before the final projection\n        context = context * gate\n        \n        # 5. Concatenate and Final Projection\n        context = context.transpose(1, 2).contiguous().view(bs, -1, self.h * self.d_k)\n        return self.out_linear(context)\n\nclass GatedTransformerLayer(nn.Module):\n    def __init__(self, d_model, num_heads, dim_feedforward=512, dropout=0.1, is_decoder=False):\n        super().__init__()\n        self.is_decoder = is_decoder\n        self.self_attn = GatedMultiHeadAttention(d_model, num_heads, dropout)\n        if is_decoder:\n            self.cross_attn = GatedMultiHeadAttention(d_model, num_heads, dropout)\n            self.norm_cross = nn.LayerNorm(d_model)\n            \n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.ffn = nn.Sequential(\n            nn.Linear(d_model, dim_feedforward),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim_feedforward, d_model)\n        )\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x, mem=None, self_mask=None, cross_mask=None):\n        # Pre-norm architecture for stability\n        # Self-Attention\n        norm_x = self.norm1(x)\n        x = x + self.dropout(self.self_attn(norm_x, norm_x, norm_x, self_mask))\n        \n        # Cross-Attention (for Decoder)\n        if self.is_decoder and mem is not None:\n            norm_x = self.norm_cross(x)\n            x = x + self.dropout(self.cross_attn(norm_x, mem, mem, cross_mask))\n            \n        # Feed-Forward\n        norm_x = self.norm2(x)\n        x = x + self.dropout(self.ffn(norm_x))\n        return x\n\nclass GatedTransformer(nn.Module):\n    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=256, num_heads=8, num_layers=3, dropout=0.1):\n        super().__init__()\n        self.encoder_emb = nn.Embedding(src_vocab_size, d_model)\n        self.decoder_emb = nn.Embedding(tgt_vocab_size, d_model)\n        self.pos_enc = PositionalEncoding(d_model)\n        \n        self.encoder_layers = nn.ModuleList([\n            GatedTransformerLayer(d_model, num_heads, dropout=dropout) for _ in range(num_layers)\n        ])\n        self.decoder_layers = nn.ModuleList([\n            GatedTransformerLayer(d_model, num_heads, dropout=dropout, is_decoder=True) for _ in range(num_layers)\n        ])\n        \n        self.fc_out = nn.Linear(d_model, tgt_vocab_size)\n\n    def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n        # Encoder\n        enc_out = self.pos_enc(self.encoder_emb(src))\n        for layer in self.encoder_layers:\n            enc_out = layer(enc_out, self_mask=src_mask)\n            \n        # Decoder\n        dec_out = self.pos_enc(self.decoder_emb(tgt))\n        # Generate triangular mask if none provided\n        if tgt_mask is None:\n            sz = tgt.size(1)\n            tgt_mask = torch.triu(torch.ones(sz, sz), diagonal=1).bool().to(src.device)\n            # Standard mask logic: 1 is masked, 0 is allowed. In our Gated MHA, we use masked_fill(mask==0, -1e9)\n            # So we invert this to: 1 is allowed, 0 is masked.\n            tgt_mask = ~tgt_mask \n\n        for layer in self.decoder_layers:\n            dec_out = layer(dec_out, mem=enc_out, self_mask=tgt_mask)\n            \n        return self.fc_out(dec_out)\n\n# ============================================================\n# 2. Data Processing (Unchanged)\n# ============================================================\n\ndef clean_ws(x: str) -> str:\n    if pd.isna(x): return \"\"\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\nclass SimpleTokenizer:\n    def __init__(self, texts):\n        self.vocab = {\"<PAD>\": 0, \"<SOS>\": 1, \"<EOS>\": 2, \"<UNK>\": 3}\n        for text in texts:\n            for char in text:\n                if char not in self.vocab:\n                    self.vocab[char] = len(self.vocab)\n        self.inv_vocab = {v: k for k, v in self.vocab.items()}\n\n    def encode(self, text, max_len):\n        ids = [self.vocab.get(c, 3) for c in text]\n        ids = [1] + ids[:max_len-2] + [2]\n        return ids + [0] * (max_len - len(ids))\n\n    def decode(self, ids):\n        return \"\".join([self.inv_vocab.get(i, \"\") for i in ids if i > 2])\n\n# ============================================================\n# 3. Training Setup (Updated to use GatedTransformer)\n# ============================================================\n\n# Assuming train.csv exists in the environment\ntrain_df = pd.read_csv(\"/kaggle/input/deep-past-initiative-machine-translation/train.csv\")\nsrc_texts = train_df['transliteration'].apply(clean_ws).tolist()\ntgt_texts = train_df['translation'].apply(clean_ws).tolist()\n\nsrc_tokenizer = SimpleTokenizer(src_texts)\ntgt_tokenizer = SimpleTokenizer(tgt_texts)\n\nclass CustomDataset(Dataset):\n    def __init__(self, src_texts, tgt_texts, src_tok, tgt_tok, max_len=128):\n        self.src_texts = src_texts\n        self.tgt_texts = tgt_texts\n        self.src_tok = src_tok\n        self.tgt_tok = tgt_tok\n        self.max_len = max_len\n\n    def __len__(self): return len(self.src_texts)\n\n    def __getitem__(self, idx):\n        return {\n            \"src\": torch.tensor(self.src_tok.encode(self.src_texts[idx], self.max_len)),\n            \"tgt\": torch.tensor(self.tgt_tok.encode(self.tgt_texts[idx], self.max_len))\n        }\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n# Using the new GatedTransformer\nmodel = GatedTransformer(len(src_tokenizer.vocab), len(tgt_tokenizer.vocab)).to(device)\noptimizer = optim.Adam(model.parameters(), lr=0.0001)\ncriterion = nn.CrossEntropyLoss(ignore_index=0)\n\n# ============================================================\n# 4. Training Loop\n# ============================================================\n\ndef train_epoch(model, loader, optimizer, criterion):\n    model.train()\n    total_loss = 0\n    for batch in loader:\n        src = batch[\"src\"].to(device)\n        tgt = batch[\"tgt\"].to(device)\n\n        tgt_input = tgt[:, :-1]\n        tgt_output = tgt[:, 1:]\n\n        optimizer.zero_grad()\n        output = model(src, tgt_input)\n        \n        loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n    return total_loss / len(loader)\n\nloader = DataLoader(CustomDataset(src_texts, tgt_texts, src_tokenizer, tgt_tokenizer), batch_size=32)\nprint(f\"Loss: {train_epoch(model, loader, optimizer, criterion)}\")\n\n# ============================================================\n# 5. Inference\n# ============================================================\n\ndef translate(text, model, src_tok, tgt_tok, max_len=128):\n    model.eval()\n    src = torch.tensor([src_tok.encode(text, max_len)]).to(device)\n    tgt_indices = [1] # <SOS>\n\n    for _ in range(max_len):\n        tgt_tensor = torch.tensor([tgt_indices]).to(device)\n        with torch.no_grad():\n            output = model(src, tgt_tensor)\n        \n        next_token = output[0, -1].argmax().item()\n        tgt_indices.append(next_token)\n        if next_token == 2: # <EOS>\n            break\n            \n    return tgt_tok.decode(tgt_indices)\n\nprint(translate(\"sample text\", model, src_tokenizer, tgt_tokenizer))","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-10T15:09:14.906547Z","iopub.execute_input":"2026-01-10T15:09:14.907136Z","iopub.status.idle":"2026-01-10T15:09:19.573271Z","shell.execute_reply.started":"2026-01-10T15:09:14.907105Z","shell.execute_reply":"2026-01-10T15:09:19.572495Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nComplete Gated Attention Training Script\nBased on arXiv:2505.06708v1 \"Gated Attention for Large Language Models\"\nAdapted for Kaggle Deep Past Initiative Machine Translation Competition\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport math\nimport pandas as pd\nimport re\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nimport os\n\n# ============================================================\n# 1. Gated Model Components (Corrected Implementation)\n# ============================================================\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super().__init__()\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() * (-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        self.register_buffer('pe', pe.unsqueeze(0))\n\n    def forward(self, x):\n        return x + self.pe[:, :x.size(1)]\n\n\nclass GatedMultiHeadAttention(nn.Module):\n    \"\"\"\n    SDPA Elementwise Gating (G1) - Most effective variant per paper.\n    Gate is computed from pre-normalized hidden states (query-dependent).\n    \"\"\"\n    def __init__(self, d_model, num_heads, dropout=0.1, gating_type='elementwise'):\n        super().__init__()\n        assert d_model % num_heads == 0\n        self.d_k = d_model // num_heads\n        self.h = num_heads\n        self.d_model = d_model\n        self.gating_type = gating_type\n        \n        # Q, K, V projections\n        self.q_linear = nn.Linear(d_model, d_model)\n        self.k_linear = nn.Linear(d_model, d_model)\n        self.v_linear = nn.Linear(d_model, d_model)\n        \n        # Gating parameters - head-specific, computed from pre-norm hidden states\n        if gating_type == 'elementwise':\n            self.gate_linear = nn.Linear(d_model, d_model)\n        elif gating_type == 'headwise':\n            self.gate_linear = nn.Linear(d_model, num_heads)\n        \n        self.out_linear = nn.Linear(d_model, d_model)\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x, k_input=None, v_input=None, mask=None):\n        \"\"\"\n        x: Pre-normalized hidden states (used for Q and gating)\n        k_input, v_input: For cross-attention (None for self-attention)\n        \"\"\"\n        bs = x.size(0)\n        seq_len = x.size(1)\n        \n        if k_input is None:\n            k_input = x\n        if v_input is None:\n            v_input = x\n        \n        # Compute gate from X (pre-normalized hidden states) - QUERY-DEPENDENT\n        if self.gating_type == 'elementwise':\n            gate = torch.sigmoid(self.gate_linear(x))\n            gate = gate.view(bs, seq_len, self.h, self.d_k).transpose(1, 2)\n        else:  # headwise\n            gate = torch.sigmoid(self.gate_linear(x))\n            gate = gate.transpose(1, 2).unsqueeze(-1)\n\n        # Q, K, V projections\n        q = self.q_linear(x).view(bs, -1, self.h, self.d_k).transpose(1, 2)\n        k = self.k_linear(k_input).view(bs, -1, self.h, self.d_k).transpose(1, 2)\n        v = self.v_linear(v_input).view(bs, -1, self.h, self.d_k).transpose(1, 2)\n\n        # Scaled Dot-Product Attention\n        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n        if mask is not None:\n            scores = scores.masked_fill(mask == 0, float('-inf'))\n        \n        attn = torch.softmax(scores, dim=-1)\n        attn = self.dropout(attn)\n        context = torch.matmul(attn, v)\n        \n        # Apply gating at G1 position (after SDPA, before concat)\n        context = context * gate\n        \n        # Concatenate and project\n        context = context.transpose(1, 2).contiguous().view(bs, -1, self.d_model)\n        return self.out_linear(context)\n\n\nclass GatedTransformerLayer(nn.Module):\n    \"\"\"Pre-norm transformer layer with gated attention.\"\"\"\n    def __init__(self, d_model, num_heads, dim_feedforward=512, dropout=0.1, \n                 is_decoder=False, gating_type='elementwise'):\n        super().__init__()\n        self.is_decoder = is_decoder\n        \n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.self_attn = GatedMultiHeadAttention(d_model, num_heads, dropout, gating_type)\n        \n        if is_decoder:\n            self.norm_cross = nn.LayerNorm(d_model)\n            self.cross_attn = GatedMultiHeadAttention(d_model, num_heads, dropout, gating_type)\n        \n        self.ffn = nn.Sequential(\n            nn.Linear(d_model, dim_feedforward),\n            nn.GELU(),  # GELU often works better than ReLU\n            nn.Dropout(dropout),\n            nn.Linear(dim_feedforward, d_model)\n        )\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x, mem=None, self_mask=None, cross_mask=None):\n        # Pre-norm self-attention\n        norm_x = self.norm1(x)\n        x = x + self.dropout(self.self_attn(norm_x, mask=self_mask))\n        \n        # Cross-attention (decoder only)\n        if self.is_decoder and mem is not None:\n            norm_x = self.norm_cross(x)\n            x = x + self.dropout(self.cross_attn(norm_x, k_input=mem, v_input=mem, mask=cross_mask))\n        \n        # Pre-norm FFN\n        norm_x = self.norm2(x)\n        x = x + self.dropout(self.ffn(norm_x))\n        return x\n\n\nclass GatedTransformer(nn.Module):\n    \"\"\"Full encoder-decoder transformer with gated attention.\"\"\"\n    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=256, num_heads=8, \n                 num_layers=4, dim_feedforward=1024, dropout=0.1, gating_type='elementwise'):\n        super().__init__()\n        self.d_model = d_model\n        \n        self.encoder_emb = nn.Embedding(src_vocab_size, d_model)\n        self.decoder_emb = nn.Embedding(tgt_vocab_size, d_model)\n        self.pos_enc = PositionalEncoding(d_model)\n        self.embed_scale = math.sqrt(d_model)\n        \n        self.encoder_layers = nn.ModuleList([\n            GatedTransformerLayer(d_model, num_heads, dim_feedforward, dropout, \n                                  is_decoder=False, gating_type=gating_type)\n            for _ in range(num_layers)\n        ])\n        \n        self.decoder_layers = nn.ModuleList([\n            GatedTransformerLayer(d_model, num_heads, dim_feedforward, dropout,\n                                  is_decoder=True, gating_type=gating_type)\n            for _ in range(num_layers)\n        ])\n        \n        self.encoder_norm = nn.LayerNorm(d_model)\n        self.decoder_norm = nn.LayerNorm(d_model)\n        self.fc_out = nn.Linear(d_model, tgt_vocab_size)\n        \n        self._init_weights()\n\n    def _init_weights(self):\n        for p in self.parameters():\n            if p.dim() > 1:\n                nn.init.xavier_uniform_(p)\n\n    def _generate_causal_mask(self, sz, device):\n        mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1)\n        mask = (mask == 0).unsqueeze(0).unsqueeze(0)\n        return mask\n\n    def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n        device = src.device\n        \n        # Encoder\n        enc_out = self.pos_enc(self.encoder_emb(src) * self.embed_scale)\n        for layer in self.encoder_layers:\n            enc_out = layer(enc_out, self_mask=src_mask)\n        enc_out = self.encoder_norm(enc_out)\n        \n        # Decoder\n        dec_out = self.pos_enc(self.decoder_emb(tgt) * self.embed_scale)\n        if tgt_mask is None:\n            tgt_mask = self._generate_causal_mask(tgt.size(1), device)\n        \n        for layer in self.decoder_layers:\n            dec_out = layer(dec_out, mem=enc_out, self_mask=tgt_mask)\n        dec_out = self.decoder_norm(dec_out)\n        \n        return self.fc_out(dec_out)\n\n\n# ============================================================\n# 2. Data Processing\n# ============================================================\n\ndef clean_ws(x: str) -> str:\n    \"\"\"Clean whitespace from text.\"\"\"\n    if pd.isna(x):\n        return \"\"\n    return re.sub(r\"\\s+\", \" \", str(x)).strip()\n\n\nclass SimpleTokenizer:\n    \"\"\"Character-level tokenizer with special tokens.\"\"\"\n    def __init__(self, texts):\n        self.vocab = {\"<PAD>\": 0, \"<SOS>\": 1, \"<EOS>\": 2, \"<UNK>\": 3}\n        for text in texts:\n            for char in text:\n                if char not in self.vocab:\n                    self.vocab[char] = len(self.vocab)\n        self.inv_vocab = {v: k for k, v in self.vocab.items()}\n        self.pad_id = 0\n        self.sos_id = 1\n        self.eos_id = 2\n        self.unk_id = 3\n\n    def encode(self, text, max_len):\n        ids = [self.vocab.get(c, self.unk_id) for c in text]\n        ids = [self.sos_id] + ids[:max_len-2] + [self.eos_id]\n        return ids + [self.pad_id] * (max_len - len(ids))\n\n    def decode(self, ids):\n        tokens = []\n        for i in ids:\n            if i == self.eos_id:\n                break\n            if i > 2:  # Skip PAD, SOS, EOS\n                tokens.append(self.inv_vocab.get(i, \"\"))\n        return \"\".join(tokens)\n\n\nclass TranslationDataset(Dataset):\n    \"\"\"Dataset for translation task.\"\"\"\n    def __init__(self, src_texts, tgt_texts, src_tok, tgt_tok, max_len=128):\n        self.src_texts = src_texts\n        self.tgt_texts = tgt_texts\n        self.src_tok = src_tok\n        self.tgt_tok = tgt_tok\n        self.max_len = max_len\n\n    def __len__(self):\n        return len(self.src_texts)\n\n    def __getitem__(self, idx):\n        return {\n            \"src\": torch.tensor(self.src_tok.encode(self.src_texts[idx], self.max_len)),\n            \"tgt\": torch.tensor(self.tgt_tok.encode(self.tgt_texts[idx], self.max_len))\n        }\n\n\n# ============================================================\n# 3. Training & Validation\n# ============================================================\n\ndef train_model(model, train_loader, val_loader, epochs, lr, device, save_path=\"best_model.pt\"):\n    \"\"\"\n    Training loop with validation and checkpointing.\n    Gated Attention enables stable training at higher learning rates.\n    \"\"\"\n    criterion = nn.CrossEntropyLoss(ignore_index=0)\n    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)\n    \n    best_val_loss = float('inf')\n\n    for epoch in range(epochs):\n        # Training\n        model.train()\n        train_loss = 0\n        train_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{epochs} [Train]\")\n        \n        for batch in train_bar:\n            src = batch[\"src\"].to(device)\n            tgt = batch[\"tgt\"].to(device)\n\n            tgt_input = tgt[:, :-1]\n            tgt_output = tgt[:, 1:]\n\n            optimizer.zero_grad()\n            output = model(src, tgt_input)\n            loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))\n            \n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n            optimizer.step()\n            \n            train_loss += loss.item()\n            train_bar.set_postfix(loss=f\"{loss.item():.4f}\")\n\n        avg_train_loss = train_loss / len(train_loader)\n\n        # Validation\n        model.eval()\n        val_loss = 0\n        val_bar = tqdm(val_loader, desc=f\"Epoch {epoch+1}/{epochs} [Val]\")\n        \n        with torch.no_grad():\n            for batch in val_bar:\n                src = batch[\"src\"].to(device)\n                tgt = batch[\"tgt\"].to(device)\n                \n                tgt_input = tgt[:, :-1]\n                tgt_output = tgt[:, 1:]\n\n                output = model(src, tgt_input)\n                loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))\n                val_loss += loss.item()\n                val_bar.set_postfix(loss=f\"{loss.item():.4f}\")\n\n        avg_val_loss = val_loss / len(val_loader)\n        scheduler.step(avg_val_loss)\n        \n        current_lr = optimizer.param_groups[0]['lr']\n        print(f\"\\nEpoch {epoch+1}: Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | LR: {current_lr:.6f}\")\n\n        # Save best model\n        if avg_val_loss < best_val_loss:\n            best_val_loss = avg_val_loss\n            torch.save({\n                'epoch': epoch,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'loss': best_val_loss,\n            }, save_path)\n            print(f\"✓ Best model saved (Val Loss: {best_val_loss:.4f})\")\n\n    return best_val_loss\n\n\ndef translate(text, model, src_tok, tgt_tok, device, max_len=128):\n    \"\"\"Greedy decoding for inference.\"\"\"\n    model.eval()\n    src = torch.tensor([src_tok.encode(text, max_len)]).to(device)\n    tgt_indices = [tgt_tok.sos_id]\n\n    for _ in range(max_len):\n        tgt_tensor = torch.tensor([tgt_indices]).to(device)\n        with torch.no_grad():\n            output = model(src, tgt_tensor)\n        \n        next_token = output[0, -1].argmax().item()\n        tgt_indices.append(next_token)\n        \n        if next_token == tgt_tok.eos_id:\n            break\n            \n    return tgt_tok.decode(tgt_indices)\n\n\ndef generate_submission(model, test_df, src_tok, tgt_tok, device, output_path=\"submission.csv\"):\n    \"\"\"Generate submission file for Kaggle.\"\"\"\n    model.eval()\n    predictions = []\n    \n    for idx, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"Generating predictions\"):\n        src_text = clean_ws(row['transliteration'])\n        translation = translate(src_text, model, src_tok, tgt_tok, device)\n        predictions.append(translation)\n    \n    submission = pd.DataFrame({\n        'id': test_df['id'],\n        'translation': predictions\n    })\n    submission.to_csv(output_path, index=False)\n    print(f\"Submission saved to {output_path}\")\n    return submission\n\n\n# ============================================================\n# 4. Main Execution\n# ============================================================\n\nif __name__ == \"__main__\":\n    # Configuration\n    CONFIG = {\n        'batch_size': 64,\n        'epochs': 30,\n        'learning_rate': 5e-4,  # Gated attention stable at higher LR\n        'd_model': 256,\n        'num_heads': 8,\n        'num_layers': 4,\n        'dim_feedforward': 1024,\n        'dropout': 0.1,\n        'max_len': 128,\n        'gating_type': 'elementwise',  # Best variant per paper\n        'train_split': 0.9,\n    }\n    \n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {DEVICE}\")\n    \n    # Load data - adjust path for Kaggle\n    DATA_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/\"\n    \n    # For local testing, use:\n    # DATA_PATH = \"./\"\n    \n    print(\"Loading data...\")\n    train_df = pd.read_csv(os.path.join(DATA_PATH, \"train.csv\"))\n    \n    # Clean text\n    src_texts = train_df['transliteration'].apply(clean_ws).tolist()\n    tgt_texts = train_df['translation'].apply(clean_ws).tolist()\n    \n    print(f\"Total samples: {len(src_texts)}\")\n    print(f\"Sample source: {src_texts[0][:50]}...\")\n    print(f\"Sample target: {tgt_texts[0][:50]}...\")\n    \n    # Build tokenizers\n    print(\"Building tokenizers...\")\n    src_tokenizer = SimpleTokenizer(src_texts)\n    tgt_tokenizer = SimpleTokenizer(tgt_texts)\n    print(f\"Source vocab size: {len(src_tokenizer.vocab)}\")\n    print(f\"Target vocab size: {len(tgt_tokenizer.vocab)}\")\n    \n    # Train/Val split\n    split_idx = int(len(src_texts) * CONFIG['train_split'])\n    \n    train_ds = TranslationDataset(\n        src_texts[:split_idx], tgt_texts[:split_idx], \n        src_tokenizer, tgt_tokenizer, CONFIG['max_len']\n    )\n    val_ds = TranslationDataset(\n        src_texts[split_idx:], tgt_texts[split_idx:], \n        src_tokenizer, tgt_tokenizer, CONFIG['max_len']\n    )\n    \n    train_loader = DataLoader(train_ds, batch_size=CONFIG['batch_size'], shuffle=True, num_workers=2)\n    val_loader = DataLoader(val_ds, batch_size=CONFIG['batch_size'], num_workers=2)\n    \n    print(f\"Train batches: {len(train_loader)} | Val batches: {len(val_loader)}\")\n    \n    # Initialize model\n    model = GatedTransformer(\n        src_vocab_size=len(src_tokenizer.vocab),\n        tgt_vocab_size=len(tgt_tokenizer.vocab),\n        d_model=CONFIG['d_model'],\n        num_heads=CONFIG['num_heads'],\n        num_layers=CONFIG['num_layers'],\n        dim_feedforward=CONFIG['dim_feedforward'],\n        dropout=CONFIG['dropout'],\n        gating_type=CONFIG['gating_type']\n    ).to(DEVICE)\n    \n    # Count parameters\n    total_params = sum(p.numel() for p in model.parameters())\n    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    print(f\"Total parameters: {total_params:,}\")\n    print(f\"Trainable parameters: {trainable_params:,}\")\n    \n    # Train\n    print(\"\\n\" + \"=\"*50)\n    print(\"Starting training with Gated Attention...\")\n    print(\"=\"*50 + \"\\n\")\n    \n    best_loss = train_model(\n        model, train_loader, val_loader, \n        CONFIG['epochs'], CONFIG['learning_rate'], \n        DEVICE, save_path=\"best_gated_model.pt\"\n    )\n    \n    print(f\"\\nTraining complete! Best validation loss: {best_loss:.4f}\")\n    \n    # Load best model for inference\n    checkpoint = torch.load(\"best_gated_model.pt\")\n    model.load_state_dict(checkpoint['model_state_dict'])\n    print(f\"Loaded best model from epoch {checkpoint['epoch']+1}\")\n    \n    # Test translation on a sample\n    print(\"\\n\" + \"=\"*50)\n    print(\"Sample Translations:\")\n    print(\"=\"*50)\n    for i in range(min(5, len(src_texts))):\n        src = src_texts[i]\n        pred = translate(src, model, src_tokenizer, tgt_tokenizer, DEVICE)\n        actual = tgt_texts[i]\n        print(f\"\\nSource: {src[:60]}...\")\n        print(f\"Predicted: {pred[:60]}...\")\n        print(f\"Actual: {actual[:60]}...\")\n    \n    # Generate submission if test.csv exists\n    test_path = os.path.join(DATA_PATH, \"test.csv\")\n    if os.path.exists(test_path):\n        print(\"\\n\" + \"=\"*50)\n        print(\"Generating submission...\")\n        print(\"=\"*50)\n        test_df = pd.read_csv(test_path)\n        generate_submission(model, test_df, src_tokenizer, tgt_tokenizer, DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T15:11:46.391834Z","iopub.execute_input":"2026-01-10T15:11:46.392132Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport pandas as pd\nimport re\nfrom tqdm import tqdm\n\n# ============================================================\n# 1. Configuration & Paths\n# ============================================================\nTEST_PATH = '/kaggle/input/deep-past-initiative-machine-translation/test.csv'\ntest_df = pd.read_csv(TEST_PATH)\n\nMAX_TGT = 128\nMAX_SRC = 128\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# ============================================================\n# 2. Post-Processing Logic\n# ============================================================\n\ndef make_one_sentence(en: str) -> str:\n    en = clean_ws(en)\n    if not en:\n        return \"...\"\n    # Split into sentences based on punctuation\n    parts = re.split(r\"(?<=[.!?])\\s+\", en)\n    parts = [p.strip() for p in parts if p.strip()]\n    # Remove ending punct from individual parts\n    parts = [re.sub(r\"[.!?]+\\s*$\", \"\", p).strip() for p in parts]\n    parts = [p for p in parts if p]\n    # Join with semicolons and ensure a single trailing period\n    out = \"; \".join(parts).strip()\n    if out and not out.endswith(\".\"):\n        out += \".\"\n    return out\n\n# ============================================================\n# 3. Gated Inference (Greedy Decoding)\n# ============================================================\n\ndef generate_translation(src_text: str, model, src_tok, tgt_tok):\n    \"\"\"\n    Manual implementation of .generate() for the custom Gated Transformer.\n    Includes the prompt formatting and greedy decoding.\n    \"\"\"\n    model.eval()\n    \n    # 1. Format the prompt\n    # Note: If the model was trained on raw transliteration -> translation, \n    # it is often better to pass just the raw text. If it was trained with \n    # instructions, keep this format.\n    prompt = f\"Instruction: Translate to English.\\nInput: {clean_ws(src_text)}\\nOutput:\"\n    \n    # 2. Tokenize\n    src_ids = src_tok.encode(prompt, MAX_SRC)\n    src_tensor = torch.tensor([src_ids]).to(DEVICE)\n    \n    # 3. Decoding Loop (Greedy)\n    # Start with <SOS> token (ID 1)\n    tgt_indices = [1] \n    \n    for _ in range(MAX_TGT):\n        tgt_tensor = torch.tensor([tgt_indices]).to(DEVICE)\n        \n        with torch.no_grad():\n            # The Gated Transformer uses its internal sigmoid gates \n            # here to stabilize the attention scores for long sequences\n            output = model(src_tensor, tgt_tensor)\n        \n        # Get the last token's distribution\n        next_token_logits = output[0, -1, :]\n        \n        # Greedy selection (you can replace this with Beam Search if needed)\n        next_token = torch.argmax(next_token_logits).item()\n        \n        tgt_indices.append(next_token)\n        \n        # Break if <EOS> (ID 2) is reached\n        if next_token == 2:\n            break\n            \n    # 4. Decode and Clean\n    raw_pred = tgt_tok.decode(tgt_indices)\n    return make_one_sentence(raw_pred)\n\n# ============================================================\n# 4. Run Inference on Test Set\n# ============================================================\n\n# Ensure model is on the correct device\nmodel.to(DEVICE)\nmodel.eval()\n\nprint(f\"Starting inference on {len(test_df)} samples...\")\npreds = []\n\n# Using tqdm for progress tracking\nfor t in tqdm(test_df[\"transliteration\"].tolist(), desc=\"Translating\"):\n    # Pass the model and tokenizers defined in your training script\n    translation = generate_translation(t, model, src_tokenizer, tgt_tokenizer)\n    preds.append(translation)\n\n# ============================================================\n# 5. Save Submission\n# ============================================================\n\nsub = pd.DataFrame({\n    \"id\": test_df[\"id\"], \n    \"translation\": preds\n})\n\nsub.to_csv(\"submission.csv\", index=False)\nprint(\"\\nSuccessfully wrote submission.csv\")\nprint(sub.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T15:11:37.806345Z","iopub.execute_input":"2026-01-10T15:11:37.806584Z","iopub.status.idle":"2026-01-10T15:11:41.759291Z","shell.execute_reply.started":"2026-01-10T15:11:37.806562Z","shell.execute_reply":"2026-01-10T15:11:41.758591Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}