{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bae751d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "from tokenizer import Tokenizer\n",
    "from model.generator import Generator\n",
    "from model.encoder import Encoder\n",
    "from model.decoder import Decoder\n",
    "from model.attn import BahdanauAttention\n",
    "from dataset import OpenMPDataset\n",
    "from accelera.src.utils.code_utils import pragma_to_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e30f61",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(vocab_size=8000)\n",
    "tokenizer.load(\"tokenizer.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db130c45",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_inputs, train_outputs = [], []\n",
    "val_inputs, val_outputs = [], []\n",
    "\n",
    "with open('../../data/data.json', 'r') as f:\n",
    "    lines = f.readlines()\n",
    "    \n",
    "    split_idx = int(0.9 * len(lines))\n",
    "    train_lines = lines[:split_idx]\n",
    "    val_lines = lines[split_idx:]\n",
    "\n",
    "for line in train_lines:\n",
    "    item = json.loads(line.strip())\n",
    "    \n",
    "    if item['label'] == 'False':\n",
    "        continue\n",
    "    \n",
    "    cls = pragma_to_class(item['label'], item['pragma'])\n",
    "    if cls == 'none':\n",
    "        continue\n",
    "    \n",
    "    input_str = f\"[CLS:{cls}] {item['code']}\"\n",
    "    output_str = item['pragma'].strip()\n",
    "    \n",
    "    if not output_str:\n",
    "        continue\n",
    "    \n",
    "    train_inputs.append(input_str)\n",
    "    train_outputs.append(output_str)\n",
    "\n",
    "for line in val_lines:\n",
    "    item = json.loads(line.strip())\n",
    "    if item['label'] == 'False':\n",
    "        continue\n",
    "    \n",
    "    cls = pragma_to_class(item['label'], item['pragma'])\n",
    "    if cls == 'none':\n",
    "        continue\n",
    "    \n",
    "    input_str = f\"[CLS:{cls}] {item['code']}\"\n",
    "    output_str = item['pragma'].strip()\n",
    "    if not output_str:\n",
    "        continue\n",
    "    \n",
    "    val_inputs.append(input_str)\n",
    "    val_outputs.append(output_str)\n",
    "\n",
    "print(f\"Training samples: {len(train_inputs)}\")\n",
    "print(f\"Validation samples: {len(val_inputs)}\")\n",
    "print(f\"\\nSample input (first 70 chars):\\n{train_inputs[0]}\")\n",
    "print(f\"Sample output:\\n{train_outputs[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5747915",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = OpenMPDataset(\n",
    "    train_inputs, train_outputs, tokenizer,\n",
    "    max_input_len=1500,\n",
    "    max_output_len=300\n",
    ")\n",
    "\n",
    "val_dataset = OpenMPDataset(\n",
    "    val_inputs, val_outputs, tokenizer,\n",
    "    max_input_len=1500,\n",
    "    max_output_len=300\n",
    ")\n",
    "\n",
    "print(f\"\\nDataset shapes:\")\n",
    "print(f\"  Train: {len(train_dataset)} samples\")\n",
    "print(f\"  Val: {len(val_dataset)} samples\")\n",
    "print(f\"  Sample input tensor shape: {train_dataset[0]['input'].shape}\")\n",
    "print(f\"  Sample output tensor shape: {train_dataset[0]['output'].shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5252d457",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=8,\n",
    "    shuffle=True,\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=8,\n",
    "    shuffle=False,\n",
    "    pin_memory=True\n",
    ")\n",
    "\n",
    "print(f\"\\n✓ Dataloaders ready!\")\n",
    "print(f\"  Train batches: {len(train_loader)}\")\n",
    "print(f\"  Val batches: {len(val_loader)}\")\n",
    "\n",
    "sample_batch = next(iter(train_loader))\n",
    "print(f\"\\nSample batch structure:\")\n",
    "print(f\"  input shape: {sample_batch['input'].shape}\")\n",
    "print(f\"  output shape: {sample_batch['output'].shape}\")\n",
    "print(f\"  input_len shape: {sample_batch['input_len'].shape}\")\n",
    "print(f\"  First sample input_len: {sample_batch['input_len'][0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11631bed",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "VOCAB_SIZE = tokenizer.vocab_size\n",
    "EMBED_SIZE = 128\n",
    "HIDDEN_SIZE = 256\n",
    "NUM_LAYERS = 3\n",
    "DROPOUT = 0.2\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "encoder = Encoder(VOCAB_SIZE, EMBED_SIZE, HIDDEN_SIZE, NUM_LAYERS, DROPOUT)\n",
    "attention = BahdanauAttention(HIDDEN_SIZE)\n",
    "decoder = Decoder(VOCAB_SIZE, EMBED_SIZE, HIDDEN_SIZE, attention, NUM_LAYERS, DROPOUT)\n",
    "model = Generator(encoder, decoder, device).to(device)\n",
    "model.apply(model._init_weights)\n",
    "\n",
    "print(\"Model architecture:\")\n",
    "print(model)\n",
    "print(f\"\\nTotal parameters: {sum(p.numel() for p in model.parameters()):,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d3125a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "PAD_IDX = tokenizer.char2idx['<PAD>']\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    optimizer, \n",
    "    mode='min', \n",
    "    factor=0.5, \n",
    "    patience=2, \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794c40e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip=1.0, teacher_forcing_ratio=0.8):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for batch in tqdm(iterator, desc=\"Training\", leave=False):\n",
    "        src = batch['input'].to(device)\n",
    "        trg = batch['output'].to(device)\n",
    "        src_len = batch['input_len'].to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(src, src_len, trg, teacher_forcing_ratio)\n",
    "        output_dim = output.shape[-1]\n",
    "        output = output[1:].view(-1, output_dim)\n",
    "        trg = trg.transpose(0, 1) \n",
    "        trg = trg[1:].reshape(-1)\n",
    "        \n",
    "        loss = criterion(output, trg)\n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        epoch_loss += loss.item()\n",
    "    \n",
    "    return epoch_loss / len(iterator)\n",
    "\n",
    "\n",
    "def evaluate(model, iterator, criterion):\n",
    "    model.eval()\n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for batch in tqdm(iterator, desc=\"Evaluating\", leave=False):\n",
    "            src = batch['input'].to(device)\n",
    "            trg = batch['output'].to(device)\n",
    "            src_len = batch['input_len'].to(device)\n",
    "            \n",
    "            output = model(src, src_len, trg, 0)\n",
    "            \n",
    "            output_dim = output.shape[-1]\n",
    "            output = output[1:].view(-1, output_dim)\n",
    "            \n",
    "            trg = trg.transpose(0, 1)\n",
    "            trg = trg[1:].reshape(-1)\n",
    "            \n",
    "            loss = criterion(output, trg)\n",
    "            epoch_loss += loss.item()\n",
    "    \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4bb0e92",
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 25\n",
    "CLIP = 1.0\n",
    "best_valid_loss = float('inf')\n",
    "training_history = {'train_loss': [], 'valid_loss': []}\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    start_time = time.time()\n",
    "    \n",
    "    tf_ratio = max(0.1, 0.5 * (0.9 ** epoch))\n",
    "    train_loss = train(model, train_loader, optimizer, criterion, CLIP, tf_ratio)\n",
    "    valid_loss = evaluate(model, val_loader, criterion)\n",
    "    scheduler.step(valid_loss)\n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save({\n",
    "            'epoch': epoch,\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'valid_loss': valid_loss,\n",
    "            'vocab_size': VOCAB_SIZE,\n",
    "            'embed_size': EMBED_SIZE,\n",
    "            'hidden_size': HIDDEN_SIZE,\n",
    "            'num_layers': NUM_LAYERS\n",
    "        }, 'best_model.pth')\n",
    "        save_status = \"✓ SAVED\"\n",
    "    else:\n",
    "        save_status = \" \"\n",
    "    \n",
    "    training_history['train_loss'].append(train_loss)\n",
    "    training_history['valid_loss'].append(valid_loss)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    epoch_mins = int((end_time - start_time) / 60)\n",
    "    epoch_secs = int((end_time - start_time) % 60)\n",
    "    \n",
    "    print(f'Epoch: {epoch+1:02}/{EPOCHS} | Time: {epoch_mins}m {epoch_secs}s | TF Ratio: {tf_ratio:.2f}')\n",
    "    print(f'\\tTrain Loss: {train_loss:.4f} | Val Loss: {valid_loss:.4f} | Best Val: {best_valid_loss:.4f} {save_status}')\n",
    "\n",
    "print(\"\\n\" + \"=\"*70)\n",
    "print(f\"✓ TRAINING COMPLETE!\")\n",
    "print(f\"Best validation loss: {best_valid_loss:.4f}\")\n",
    "print(f\"Model saved to 'best_model.pth'\")\n",
    "print(\"=\"*70)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6d9a8e25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded checkpoint from best_model.pth (epoch 8)\n",
      "Sample input (truncated): [CLS:reduction] for (i = 0; i < 1000; ++i)\n",
      "{\n",
      "  logic_and = logic_and && logics[i];\n",
      "}\n",
      "\n",
      "Reference pragma: omp parallel for schedule(dynamic,1) private(i) reduction(&&:logic_and)\n",
      "Greedy prediction: omp parallel for schedule(dynamic,1) private(i) reduction(&&:logic_and)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import sys\n",
    "import pathlib\n",
    "sys.path.append(str(pathlib.Path().resolve()))  # ensure local modules are importable\n",
    "import os\n",
    "\n",
    "checkpoint_path = \"best_model.pth\"\n",
    "if not os.path.exists(checkpoint_path):\n",
    "    raise FileNotFoundError(\"Run training first so 'best_model.pth' exists.\")\n",
    "\n",
    "checkpoint = torch.load(checkpoint_path, map_location=device)\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "model.eval()\n",
    "print(f\"Loaded checkpoint from {checkpoint_path} (epoch {checkpoint.get('epoch', '?')})\")\n",
    "\n",
    "SOS_IDX = tokenizer.char2idx['<SOS>']\n",
    "EOS_IDX = tokenizer.char2idx['<EOS>']\n",
    "\n",
    "# Greedy baseline (kept for comparison)\n",
    "def greedy_generate(code_snippet: str, cls: str = \"parallel\", max_len: int = 80) -> str:\n",
    "    model.eval()\n",
    "    text = code_snippet if code_snippet.startswith(\"[CLS:\") else f\"[CLS:{cls}] {code_snippet}\"\n",
    "    input_ids = tokenizer.encode(text, max_length=500, add_special_tokens=True)\n",
    "    input_len = next((i for i, tok in enumerate(input_ids) if tok == PAD_IDX), len(input_ids))\n",
    "    input_tensor = torch.tensor([input_ids], device=device)\n",
    "    input_len_tensor = torch.tensor([input_len], device=device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        enc_outs, hidden, cell = model.encoder(input_tensor, input_len_tensor)\n",
    "        mask = (torch.arange(enc_outs.size(1), device=device).unsqueeze(0) < input_len_tensor.unsqueeze(1)).float()\n",
    "\n",
    "        hidden = hidden.view(model.encoder.num_layers, 2, 1, model.encoder.hidden_size)\n",
    "        hidden = torch.cat((hidden[:, 0], hidden[:, 1]), dim=2)\n",
    "        hidden = model.hidden_projection(hidden)\n",
    "\n",
    "        cell = cell.view(model.encoder.num_layers, 2, 1, model.encoder.hidden_size)\n",
    "        cell = torch.cat((cell[:, 0], cell[:, 1]), dim=2)\n",
    "        cell = model.cell_projection(cell)\n",
    "\n",
    "        input_token = torch.tensor([SOS_IDX], device=device)\n",
    "        generated = []\n",
    "        for _ in range(max_len):\n",
    "            output, hidden, cell, _ = model.decoder(input_token, hidden, cell, enc_outs, mask)\n",
    "            top1 = output.argmax(1)\n",
    "            token_id = top1.item()\n",
    "            if token_id == EOS_IDX:\n",
    "                break\n",
    "            generated.append(token_id)\n",
    "            input_token = top1\n",
    "\n",
    "    return tokenizer.decode(generated)\n",
    "\n",
    "\n",
    "\n",
    "# Quick sanity check on a validation example\n",
    "sample_input = val_inputs[18]\n",
    "reference = val_outputs[18]\n",
    "prediction_greedy = greedy_generate(sample_input)\n",
    "print(\"Sample input (truncated):\", sample_input[:140] + \"...\" if len(sample_input) > 140 else sample_input)\n",
    "print(\"Reference pragma:\", reference)\n",
    "print(\"Greedy prediction:\", prediction_greedy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}