{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "gpuType": "T4" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "collapsed": true, "id": "12APLOKE15uD", "outputId": "fb61078b-a249-476a-af53-e43ca978c8c1" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Requirement already satisfied: torch in /usr/local/lib/python3.11/dist-packages (2.5.1+cu124)\n", "Requirement already satisfied: tqdm in /usr/local/lib/python3.11/dist-packages (4.67.1)\n", "Requirement already satisfied: streamlit in /usr/local/lib/python3.11/dist-packages (1.42.2)\n", "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch) (3.17.0)\n", "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.11/dist-packages (from torch) (4.12.2)\n", "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch) (3.4.2)\n", "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch) (3.1.5)\n", "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch) (2024.10.0)\n", "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch) (12.4.127)\n", "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch) (12.4.127)\n", "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch) (12.4.127)\n", "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.11/dist-packages (from torch) (9.1.0.70)\n", "Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /usr/local/lib/python3.11/dist-packages (from torch) (12.4.5.8)\n", "Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /usr/local/lib/python3.11/dist-packages (from torch) (11.2.1.3)\n", "Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /usr/local/lib/python3.11/dist-packages (from torch) (10.3.5.147)\n", "Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /usr/local/lib/python3.11/dist-packages (from torch) (11.6.1.9)\n", "Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /usr/local/lib/python3.11/dist-packages (from torch) (12.3.1.170)\n", "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch) (2.21.5)\n", "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch) (12.4.127)\n", "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch) (12.4.127)\n", "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.11/dist-packages (from torch) (3.1.0)\n", "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch) (1.13.1)\n", "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch) (1.3.0)\n", "Requirement already satisfied: altair<6,>=4.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (5.5.0)\n", "Requirement already satisfied: blinker<2,>=1.0.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (1.9.0)\n", "Requirement already satisfied: cachetools<6,>=4.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (5.5.1)\n", "Requirement already satisfied: click<9,>=7.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (8.1.8)\n", "Requirement already satisfied: numpy<3,>=1.23 in /usr/local/lib/python3.11/dist-packages (from streamlit) (1.26.4)\n", "Requirement already satisfied: packaging<25,>=20 in /usr/local/lib/python3.11/dist-packages (from streamlit) (24.2)\n", "Requirement already satisfied: pandas<3,>=1.4.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (2.2.2)\n", "Requirement already satisfied: pillow<12,>=7.1.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (11.1.0)\n", "Requirement already satisfied: protobuf<6,>=3.20 in /usr/local/lib/python3.11/dist-packages (from streamlit) (4.25.6)\n", "Requirement already satisfied: pyarrow>=7.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (17.0.0)\n", "Requirement already satisfied: requests<3,>=2.27 in /usr/local/lib/python3.11/dist-packages (from streamlit) (2.32.3)\n", "Requirement already satisfied: rich<14,>=10.14.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (13.9.4)\n", "Requirement already satisfied: tenacity<10,>=8.1.0 in /usr/local/lib/python3.11/dist-packages (from streamlit) (9.0.0)\n", "Requirement already satisfied: toml<2,>=0.10.1 in /usr/local/lib/python3.11/dist-packages (from streamlit) (0.10.2)\n", "Requirement already satisfied: watchdog<7,>=2.1.5 in /usr/local/lib/python3.11/dist-packages (from streamlit) (6.0.0)\n", "Requirement already satisfied: gitpython!=3.1.19,<4,>=3.0.7 in /usr/local/lib/python3.11/dist-packages (from streamlit) (3.1.44)\n", "Requirement already satisfied: pydeck<1,>=0.8.0b4 in /usr/local/lib/python3.11/dist-packages (from streamlit) (0.9.1)\n", "Requirement already satisfied: tornado<7,>=6.0.3 in /usr/local/lib/python3.11/dist-packages (from streamlit) (6.4.2)\n", "Requirement already satisfied: jsonschema>=3.0 in /usr/local/lib/python3.11/dist-packages (from altair<6,>=4.0->streamlit) (4.23.0)\n", "Requirement already satisfied: narwhals>=1.14.2 in /usr/local/lib/python3.11/dist-packages (from altair<6,>=4.0->streamlit) (1.27.1)\n", "Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.11/dist-packages (from gitpython!=3.1.19,<4,>=3.0.7->streamlit) (4.0.12)\n", "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas<3,>=1.4.0->streamlit) (2.8.2)\n", "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas<3,>=1.4.0->streamlit) (2025.1)\n", "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas<3,>=1.4.0->streamlit) (2025.1)\n", "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch) (3.0.2)\n", "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.27->streamlit) (3.4.1)\n", "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.27->streamlit) (3.10)\n", "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.27->streamlit) (2.3.0)\n", "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.27->streamlit) (2025.1.31)\n", "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.11/dist-packages (from rich<14,>=10.14.0->streamlit) (3.0.0)\n", "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.11/dist-packages (from rich<14,>=10.14.0->streamlit) (2.18.0)\n", "Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.11/dist-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.19,<4,>=3.0.7->streamlit) (5.0.2)\n", "Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (25.1.0)\n", "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (2024.10.1)\n", "Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (0.36.2)\n", "Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit) (0.22.3)\n", "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.11/dist-packages (from markdown-it-py>=2.2.0->rich<14,>=10.14.0->streamlit) (0.1.2)\n", "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/dist-packages (from python-dateutil>=2.8.2->pandas<3,>=1.4.0->streamlit) (1.17.0)\n" ] } ], "source": [ "!pip install torch tqdm streamlit" ] }, { "cell_type": "code", "source": [ "######################################\n", "# Pseudocode2Cpp.py\n", "######################################\n", "import os\n", "import streamlit as st\n", "import torch\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "import math\n", "import re\n", "from tqdm import tqdm\n", "from typing import List, Tuple\n", "import random\n", "import requests\n", "from torch.utils.data import DataLoader, TensorDataset" ], "metadata": { "id": "tEYW8hGR19sm" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# ----------------------------\n", "# 1. Hyperparameters\n", "# ----------------------------\n", "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "MAX_LEN = 128 # maximum sequence length\n", "EMBED_DIM = 256 # embedding dimension\n", "FF_DIM = 512 # feedforward dimension in Transformer\n", "NHEAD = 4 # number of heads in multihead attention\n", "NUM_ENCODER_LAYERS = 2\n", "NUM_DECODER_LAYERS = 2\n", "BATCH_SIZE = 64\n", "EPOCHS = 10 # Increase for real training\n", "LEARNING_RATE = 1e-4\n", "\n", "# Special tokens\n", "PAD_TOKEN = \"\"\n", "SOS_TOKEN = \"\"\n", "EOS_TOKEN = \"\"\n", "UNK_TOKEN = \"\"" ], "metadata": { "id": "HelkrJ-01-2B" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# ----------------------------\n", "# 2. Data Loading & Preprocessing\n", "# ----------------------------\n", "\n", "def load_spoc_data(file_path: str):\n", " \"\"\"\n", " Loads (pseudo_code, cpp_code) pairs from a TSV file or raw GitHub link.\n", " Each line is assumed to have: pseudocode c++ code.\n", " \"\"\"\n", " pairs = []\n", "\n", " # If file_path is a URL, fetch it with requests\n", " if file_path.startswith(\"http\"):\n", " response = requests.get(file_path)\n", " response.raise_for_status()\n", " lines = response.text.strip().split(\"\\n\")\n", " else:\n", " # Otherwise, assume it's a local file path\n", " with open(file_path, 'r', encoding='utf-8') as f:\n", " lines = f.readlines()\n", "\n", " for line in lines:\n", " line = line.strip()\n", " if not line:\n", " continue\n", " cols = line.split('\\t')\n", " if len(cols) >= 2:\n", " pseudo = cols[0].strip()\n", " cpp = cols[1].strip()\n", " pairs.append((pseudo, cpp))\n", "\n", " return pairs\n", "\n", "def create_dataloader(pairs, src_stoi, tgt_stoi, batch_size):\n", " src_batches = []\n", " tgt_batches = []\n", " for pseudo, cpp in pairs:\n", " src_ids = pad_sequence(numericalize(pseudo, src_stoi), MAX_LEN, src_stoi[PAD_TOKEN])\n", " tgt_ids = pad_sequence(numericalize(cpp, tgt_stoi), MAX_LEN, tgt_stoi[PAD_TOKEN])\n", " src_batches.append(src_ids)\n", " tgt_batches.append(tgt_ids)\n", "\n", " src_tensor = torch.tensor(src_batches, dtype=torch.long)\n", " tgt_tensor = torch.tensor(tgt_batches, dtype=torch.long)\n", " dataset = TensorDataset(src_tensor, tgt_tensor)\n", " return DataLoader(dataset, batch_size=batch_size, shuffle=True, pin_memory=True)\n", "\n", "def tokenize_line(text: str) -> List[str]:\n", " \"\"\"Enhanced tokenizer for pseudocode/C++ patterns\"\"\"\n", " # Separate operators and punctuation\n", " text = re.sub(r'([=+\\-*/%<>!&|^~])', r' \\1 ', text) # Operators\n", " text = re.sub(r'(?!&|^~]+|[:;{},()\\[\\]\\.]', text)\n", "\n", "def build_vocab(pairs: List[Tuple[str, str]]) -> Tuple[dict, dict, dict, dict]:\n", " \"\"\"\n", " Build source (pseudo) and target (cpp) vocabularies from training data.\n", " Returns:\n", " src_stoi, src_itos, tgt_stoi, tgt_itos\n", " \"\"\"\n", " src_words = set()\n", " tgt_words = set()\n", "\n", " for (pseudo, cpp) in pairs:\n", " for tok in tokenize_line(pseudo):\n", " src_words.add(tok)\n", " for tok in tokenize_line(cpp):\n", " tgt_words.add(tok)\n", "\n", " # Add special tokens\n", " src_vocab = [PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN] + sorted(list(src_words))\n", " tgt_vocab = [PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN] + sorted(list(tgt_words))\n", "\n", " src_stoi = {w: i for i, w in enumerate(src_vocab)}\n", " src_itos = {i: w for i, w in enumerate(src_vocab)}\n", " tgt_stoi = {w: i for i, w in enumerate(tgt_vocab)}\n", " tgt_itos = {i: w for i, w in enumerate(tgt_vocab)}\n", "\n", " return src_stoi, src_itos, tgt_stoi, tgt_itos\n", "\n", "def numericalize(text: str, stoi: dict) -> List[int]:\n", " \"\"\"\n", " Convert text string to a list of token IDs.\n", " \"\"\"\n", " tokens = tokenize_line(text)\n", " ids = []\n", " for t in tokens:\n", " if t in stoi:\n", " ids.append(stoi[t])\n", " else:\n", " ids.append(stoi[UNK_TOKEN])\n", " return ids\n", "\n", "def pad_sequence(seq: List[int], max_len: int, pad_id: int) -> List[int]:\n", " \"\"\"Proper padding with SOS/EOS handling\"\"\"\n", " seq = seq[:max_len-2] # Leave space for SOS/EOS\n", " seq = [src_stoi[SOS_TOKEN]] + seq + [src_stoi[EOS_TOKEN]] # Add control tokens\n", " padding = [pad_id] * (max_len - len(seq))\n", " return seq + padding\n", "\n", "def create_batches(pairs, src_stoi, tgt_stoi, batch_size):\n", " \"\"\"\n", " Yield batches of data (source_ids, target_ids).\n", " \"\"\"\n", " random.shuffle(pairs)\n", " for i in range(0, len(pairs), batch_size):\n", " batch_pairs = pairs[i:i+batch_size]\n", " src_batch = []\n", " tgt_batch = []\n", " for pseudo, cpp in batch_pairs:\n", " src_ids = numericalize(pseudo, src_stoi)\n", " tgt_ids = numericalize(cpp, tgt_stoi)\n", "\n", " # Pad/truncate\n", " src_ids = pad_sequence(src_ids, MAX_LEN, src_stoi[PAD_TOKEN])\n", " tgt_ids = pad_sequence(tgt_ids, MAX_LEN, tgt_stoi[PAD_TOKEN])\n", "\n", " src_batch.append(src_ids)\n", " tgt_batch.append(tgt_ids)\n", "\n", " src_batch = torch.tensor(src_batch, dtype=torch.long, device=DEVICE)\n", " tgt_batch = torch.tensor(tgt_batch, dtype=torch.long, device=DEVICE)\n", " yield src_batch, tgt_batch" ], "metadata": { "id": "2lFlkj-t2AGg" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# ----------------------------\n", "# 3. Transformer Model Implementation (from scratch)\n", "# ----------------------------\n", "\n", "class PositionalEncoding(nn.Module):\n", " def __init__(self, d_model, max_len=5000):\n", " super(PositionalEncoding, self).__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", " pe = pe.unsqueeze(0) # shape (1, max_len, d_model)\n", " self.register_buffer('pe', pe)\n", "\n", " def forward(self, x):\n", " # x shape: (batch_size, seq_len, d_model)\n", " seq_len = x.size(1)\n", " x = x + self.pe[:, :seq_len, :]\n", " return x\n", "\n", "class MultiHeadAttention(nn.Module):\n", " def __init__(self, d_model, n_heads):\n", " super(MultiHeadAttention, self).__init__()\n", " assert d_model % n_heads == 0\n", " self.d_model = d_model\n", " self.n_heads = n_heads\n", " self.head_dim = d_model // n_heads\n", "\n", " self.query_linear = nn.Linear(d_model, d_model)\n", " self.key_linear = nn.Linear(d_model, d_model)\n", " self.value_linear = nn.Linear(d_model, d_model)\n", " self.out_linear = nn.Linear(d_model, d_model)\n", "\n", " def forward(self, query, key, value, mask=None):\n", " # query/key/value shape: (batch_size, seq_len, d_model)\n", " B, Q_len, _ = query.size()\n", " B, K_len, _ = key.size()\n", " B, V_len, _ = value.size()\n", "\n", " # Linear projections\n", " Q = self.query_linear(query) # (B, Q_len, d_model)\n", " K = self.key_linear(key) # (B, K_len, d_model)\n", " V = self.value_linear(value) # (B, V_len, d_model)\n", "\n", " # Reshape for multi-head\n", " Q = Q.view(B, Q_len, self.n_heads, self.head_dim).transpose(1,2) # (B, n_heads, Q_len, head_dim)\n", " K = K.view(B, K_len, self.n_heads, self.head_dim).transpose(1,2) # (B, n_heads, K_len, head_dim)\n", " V = V.view(B, V_len, self.n_heads, self.head_dim).transpose(1,2) # (B, n_heads, V_len, head_dim)\n", "\n", " # Scaled dot-product attention\n", " scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) # (B, n_heads, Q_len, K_len)\n", " if mask is not None:\n", " scores = scores.masked_fill(mask == 0, float('-inf'))\n", " attn = torch.softmax(scores, dim=-1) # (B, n_heads, Q_len, K_len)\n", "\n", " context = torch.matmul(attn, V) # (B, n_heads, Q_len, head_dim)\n", " context = context.transpose(1,2).contiguous().view(B, Q_len, self.d_model)\n", " out = self.out_linear(context)\n", " return out\n", "\n", "class FeedForward(nn.Module):\n", " def __init__(self, d_model, dim_feedforward):\n", " super(FeedForward, self).__init__()\n", " self.fc1 = nn.Linear(d_model, dim_feedforward)\n", " self.fc2 = nn.Linear(dim_feedforward, d_model)\n", " self.relu = nn.ReLU()\n", "\n", " def forward(self, x):\n", " return self.fc2(self.relu(self.fc1(x)))\n", "\n", "class EncoderLayer(nn.Module):\n", " def __init__(self, d_model, n_heads, dim_feedforward):\n", " super(EncoderLayer, self).__init__()\n", " self.self_attn = MultiHeadAttention(d_model, n_heads)\n", " self.ff = FeedForward(d_model, dim_feedforward)\n", " self.norm1 = nn.LayerNorm(d_model)\n", " self.norm2 = nn.LayerNorm(d_model)\n", " self.dropout = nn.Dropout(0.1)\n", "\n", " def forward(self, src, src_mask=None):\n", " # Self-attention\n", " attn_out = self.self_attn(src, src, src, mask=src_mask)\n", " src = self.norm1(src + self.dropout(attn_out))\n", " # Feed Forward\n", " ff_out = self.ff(src)\n", " src = self.norm2(src + self.dropout(ff_out))\n", " return src\n", "\n", "class DecoderLayer(nn.Module):\n", " def __init__(self, d_model, n_heads, dim_feedforward):\n", " super(DecoderLayer, self).__init__()\n", " self.self_attn = MultiHeadAttention(d_model, n_heads)\n", " self.cross_attn = MultiHeadAttention(d_model, n_heads)\n", " self.ff = FeedForward(d_model, dim_feedforward)\n", "\n", " self.norm1 = nn.LayerNorm(d_model)\n", " self.norm2 = nn.LayerNorm(d_model)\n", " self.norm3 = nn.LayerNorm(d_model)\n", " self.dropout = nn.Dropout(0.1)\n", "\n", " def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):\n", " # Self-attention (mask future tokens)\n", " _tgt = tgt\n", " tgt = self.norm1(tgt + self.dropout(self.self_attn(tgt, tgt, tgt, mask=tgt_mask)))\n", " # Cross-attention\n", " _tgt2 = tgt\n", " tgt = self.norm2(tgt + self.dropout(self.cross_attn(tgt, memory, memory, mask=memory_mask)))\n", " # Feed Forward\n", " ff_out = self.ff(tgt)\n", " tgt = self.norm3(tgt + self.dropout(ff_out))\n", " return tgt\n", "\n", "class Encoder(nn.Module):\n", " def __init__(self, vocab_size, d_model, n_heads, num_layers, dim_feedforward):\n", " super(Encoder, self).__init__()\n", " self.embedding = nn.Embedding(vocab_size, d_model)\n", " self.pos_encoding = PositionalEncoding(d_model)\n", " self.layers = nn.ModuleList([\n", " EncoderLayer(d_model, n_heads, dim_feedforward)\n", " for _ in range(num_layers)\n", " ])\n", "\n", " def forward(self, src, src_mask=None):\n", " # src shape: (batch_size, seq_len)\n", " x = self.embedding(src) # (batch_size, seq_len, d_model)\n", " x = self.pos_encoding(x)\n", " for layer in self.layers:\n", " x = layer(x, src_mask)\n", " return x\n", "\n", "class Decoder(nn.Module):\n", " def __init__(self, vocab_size, d_model, n_heads, num_layers, dim_feedforward):\n", " super(Decoder, self).__init__()\n", " self.embedding = nn.Embedding(vocab_size, d_model)\n", " self.pos_encoding = PositionalEncoding(d_model)\n", " self.layers = nn.ModuleList([\n", " DecoderLayer(d_model, n_heads, dim_feedforward)\n", " for _ in range(num_layers)\n", " ])\n", " self.fc_out = nn.Linear(d_model, vocab_size)\n", "\n", " def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):\n", " x = self.embedding(tgt)\n", " x = self.pos_encoding(x)\n", " for layer in self.layers:\n", " x = layer(x, memory, tgt_mask, memory_mask)\n", " logits = self.fc_out(x) # (batch_size, seq_len, vocab_size)\n", " return logits\n", "\n", "class TransformerSeq2Seq(nn.Module):\n", " def __init__(self, src_vocab_size, tgt_vocab_size, d_model, n_heads, num_encoder_layers,\n", " num_decoder_layers, dim_feedforward):\n", " super(TransformerSeq2Seq, self).__init__()\n", " self.encoder = Encoder(src_vocab_size, d_model, n_heads, num_encoder_layers, dim_feedforward)\n", " self.decoder = Decoder(tgt_vocab_size, d_model, n_heads, num_decoder_layers, dim_feedforward)\n", "\n", " def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n", " # src: (batch_size, src_seq_len)\n", " # tgt: (batch_size, tgt_seq_len)\n", " memory = self.encoder(src, src_mask) # (batch_size, src_seq_len, d_model)\n", " outputs = self.decoder(tgt, memory, tgt_mask) # (batch_size, tgt_seq_len, vocab_size)\n", " return outputs" ], "metadata": { "id": "f8HioKcS2ZRy" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# ----------------------------\n", "# 4. Training Setup\n", "# ----------------------------\n", "import torch\n", "import torch.nn as nn\n", "from torch.utils.data import DataLoader, TensorDataset\n", "from typing import List, Tuple\n", "import random\n", "def generate_subsequent_mask(size):\n", " # Mask out subsequent positions (for decoding)\n", " mask = torch.triu(torch.ones(size, size), diagonal=1).bool()\n", " return ~mask # True where we can attend, False where we cannot\n", "\n", "def train_one_epoch(model, optimizer, criterion, train_data, src_stoi, tgt_stoi):\n", " model.train()\n", " total_loss = 0\n", " steps = 0\n", "\n", " data_loader = create_dataloader(train_pairs, src_stoi, tgt_stoi, BATCH_SIZE)\n", " for src_batch, tgt_batch in data_loader:\n", " src_batch = src_batch.to(DEVICE)\n", " tgt_batch = tgt_batch.to(DEVICE)\n", "\n", " # Prepare the target inputs and outputs (shifted by one token)\n", " tgt_inp = tgt_batch[:, :-1]\n", " tgt_out = tgt_batch[:, 1:]\n", "\n", " # Create subsequent mask for the target sequence\n", " tgt_seq_len = tgt_inp.size(1)\n", " tgt_mask = generate_subsequent_mask(tgt_seq_len).to(DEVICE)\n", "\n", " optimizer.zero_grad()\n", " logits = model(src_batch, tgt_inp, None, tgt_mask) # (B, seq_len, vocab_size)\n", "\n", " # Use .reshape() instead of .view() to avoid runtime errors\n", " loss = criterion(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1))\n", " loss.backward()\n", " optimizer.step()\n", "\n", " total_loss += loss.item()\n", " steps += 1\n", "\n", " return total_loss / steps\n", "\n", "def evaluate(model, criterion, eval_data, src_stoi, tgt_stoi):\n", " model.eval()\n", " total_loss = 0\n", " steps = 0\n", " with torch.no_grad():\n", " for src_batch, tgt_batch in create_batches(eval_data, src_stoi, tgt_stoi, BATCH_SIZE):\n", " tgt_inp = tgt_batch[:, :-1]\n", " tgt_out = tgt_batch[:, 1:]\n", " tgt_seq_len = tgt_inp.size(1)\n", " tgt_mask = generate_subsequent_mask(tgt_seq_len).to(DEVICE)\n", "\n", " logits = model(src_batch, tgt_inp, None, tgt_mask)\n", " # Use .reshape() instead of .view()\n", " loss = criterion(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1))\n", "\n", " total_loss += loss.item()\n", " steps += 1\n", " return total_loss / steps\n", "\n", "def greedy_decode(model, src, src_stoi, tgt_stoi, tgt_itos, max_len=MAX_LEN):\n", " \"\"\"\n", " Given a single source sequence (1D list of token IDs),\n", " generate a decoded target sequence using greedy search.\n", " \"\"\"\n", " model.eval()\n", " src = torch.tensor(src, dtype=torch.long, device=DEVICE).unsqueeze(0) # (1, seq_len)\n", " memory = model.encoder(src) # (1, seq_len, d_model)\n", "\n", " ys = torch.tensor([tgt_stoi[SOS_TOKEN]], dtype=torch.long, device=DEVICE).unsqueeze(0) # (1, 1)\n", " for i in range(max_len-1):\n", " tgt_mask = generate_subsequent_mask(ys.size(1)).to(DEVICE)\n", " out = model.decoder(ys, memory, tgt_mask) # (1, seq_len, vocab_size)\n", " prob = out[:, -1, :] # last timestep\n", " next_token = torch.argmax(prob, dim=1).item()\n", " ys = torch.cat([ys, torch.tensor([[next_token]], device=DEVICE)], dim=1)\n", " if next_token == tgt_stoi[EOS_TOKEN]:\n", " break\n", "\n", " # Convert back to tokens\n", " out_tokens = ys.squeeze(0).tolist() # e.g. [SOS, ..., EOS]\n", " # Remove the initial SOS\n", " out_tokens = out_tokens[1:]\n", " # Stop at EOS if present\n", " if tgt_stoi[EOS_TOKEN] in out_tokens:\n", " eos_idx = out_tokens.index(tgt_stoi[EOS_TOKEN])\n", " out_tokens = out_tokens[:eos_idx]\n", "\n", " return \" \".join(tgt_itos[t] for t in out_tokens)" ], "metadata": { "id": "ffYgGSXy2a4B" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# ----------------------------\n", "# 5. Main: Train the Model\n", "# ----------------------------\n", "if __name__ == \"__main__\":\n", " # Hardcode the file paths from your GitHub repo (raw URLs):\n", " train_path = \"https://raw.githubusercontent.com/asadsandhu/Pseudocode2Cpp/main/spoc/train/spoc-train.tsv\"\n", " eval_path = \"https://raw.githubusercontent.com/asadsandhu/Pseudocode2Cpp/main/spoc/train/split/spoc-train-eval.tsv\"\n", "\n", " print(f\"Loading training data from {train_path} ...\")\n", " train_pairs = load_spoc_data(train_path)\n", " print(f\"Loaded {len(train_pairs)} training pairs.\")\n", "\n", " print(f\"Loading eval data from {eval_path} ...\")\n", " eval_pairs = load_spoc_data(eval_path)\n", " print(f\"Loaded {len(eval_pairs)} eval pairs.\")\n", "\n", " print(\"Building vocab...\")\n", " src_stoi, src_itos, tgt_stoi, tgt_itos = build_vocab(train_pairs)\n", " global stoi_eos\n", " stoi_eos = tgt_stoi[EOS_TOKEN] # for pad_sequence usage\n", "\n", " print(\"Creating model...\")\n", " model = TransformerSeq2Seq(\n", " src_vocab_size=len(src_stoi),\n", " tgt_vocab_size=len(tgt_stoi),\n", " d_model=EMBED_DIM,\n", " n_heads=NHEAD,\n", " num_encoder_layers=NUM_ENCODER_LAYERS,\n", " num_decoder_layers=NUM_DECODER_LAYERS,\n", " dim_feedforward=FF_DIM\n", " ).to(DEVICE)\n", "\n", " criterion = nn.CrossEntropyLoss(ignore_index=tgt_stoi[PAD_TOKEN])\n", " optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n", "\n", " print(\"Starting training...\")\n", " for epoch in range(1, EPOCHS+1):\n", " train_loss = train_one_epoch(model, optimizer, criterion, train_pairs, src_stoi, tgt_stoi)\n", " eval_loss = evaluate(model, criterion, eval_pairs, src_stoi, tgt_stoi)\n", " print(f\"Epoch [{epoch}/{EPOCHS}] - Train Loss: {train_loss:.4f}, Eval Loss: {eval_loss:.4f}\")\n", "\n", " # Save model & vocab\n", " torch.save({\n", " 'model_state_dict': model.state_dict(),\n", " 'src_stoi': src_stoi,\n", " 'src_itos': src_itos,\n", " 'tgt_stoi': tgt_stoi,\n", " 'tgt_itos': tgt_itos\n", " }, \"model.pth\")\n", "\n", " print(\"Model and vocab saved to model.pth\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "iffrMhkc2cVt", "outputId": "38839989-38e5-4b10-fbea-90767dca60e3" }, "execution_count": null, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Loading training data from https://raw.githubusercontent.com/asadsandhu/Pseudocode2Cpp/main/spoc/train/spoc-train.tsv ...\n", "Loaded 293855 training pairs.\n", "Loading eval data from https://raw.githubusercontent.com/asadsandhu/Pseudocode2Cpp/main/spoc/train/split/spoc-train-eval.tsv ...\n", "Loaded 27289 eval pairs.\n", "Building vocab...\n", "Creating model...\n", "Starting training...\n", "Epoch [1/10] - Train Loss: 0.9915, Eval Loss: 0.4901\n", "Epoch [2/10] - Train Loss: 0.4401, Eval Loss: 0.3597\n", "Epoch [3/10] - Train Loss: 0.3326, Eval Loss: 0.2897\n", "Epoch [4/10] - Train Loss: 0.2752, Eval Loss: 0.2735\n", "Epoch [5/10] - Train Loss: 0.2401, Eval Loss: 0.2281\n", "Epoch [6/10] - Train Loss: 0.2166, Eval Loss: 0.2111\n", "Epoch [7/10] - Train Loss: 0.2002, Eval Loss: 0.2015\n", "Epoch [8/10] - Train Loss: 0.1883, Eval Loss: 0.1919\n", "Epoch [9/10] - Train Loss: 0.1793, Eval Loss: 0.1848\n", "Epoch [10/10] - Train Loss: 0.1724, Eval Loss: 0.1819\n", "Model and vocab saved to transformer_spoc.pth\n" ] } ] } ] }