{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":12863,"databundleVersionId":788719,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-13T10:58:20.001680Z","iopub.execute_input":"2025-12-13T10:58:20.001838Z","iopub.status.idle":"2025-12-13T10:58:21.566137Z","shell.execute_reply.started":"2025-12-13T10:58:20.001823Z","shell.execute_reply":"2025-12-13T10:58:21.565342Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport json\nimport time\nimport re\nimport random\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\n\n# ============================================================\n# Config (Kaggle paths + hyperparams)\n# ============================================================\n\nDATA_DIR = \"/kaggle/input/tensorflow2-question-answering\"\nTRAIN_PATH = os.path.join(DATA_DIR, \"simplified-nq-train.jsonl\")\nTEST_PATH = os.path.join(DATA_DIR, \"simplified-nq-test.jsonl\")\nSAMPLE_SUB_PATH = os.path.join(DATA_DIR, \"sample_submission.csv\")\n\nOUTPUT_PATH = \"/kaggle/working/submission.csv\"\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\n# Tune these if you want – they’re chosen to stay under time limits.\nif device.type == \"cuda\":\n    MAX_TRAIN_SAMPLES = 60000   # jsonl lines to read\n    EPOCHS = 4\n    MAX_NEG_PER_Q = 16          # max negatives per question\nelse:\n    MAX_TRAIN_SAMPLES = 15000   # smaller for CPU\n    EPOCHS = 3\n    MAX_NEG_PER_Q = 10\n\nMAX_Q_LEN = 32\nMAX_P_LEN = 128\nEMBED_DIM = 128\nHIDDEN_DIM = 64\nBATCH_SIZE = 64\nLR = 1e-3\nVAL_FRACTION = 0.2\n\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif device.type == \"cuda\":\n    torch.cuda.manual_seed_all(SEED)\n\n\n# ============================================================\n# Utils\n# ============================================================\n\ndef read_jsonl(path, max_lines=None):\n    count = 0\n    with open(path, \"r\", encoding=\"utf-8\") as f:\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            yield json.loads(line)\n            count += 1\n            if max_lines is not None and count >= max_lines:\n                break\n\n\nTOKEN_PATTERN = re.compile(r\"\\w+|[^\\w\\s]\", re.UNICODE)\n\n\ndef tokenize(text):\n    return TOKEN_PATTERN.findall(text.lower())\n\n\ndef tokens_to_ids(tokens, vocab, max_len):\n    unk = vocab[\"<UNK>\"]\n    pad = vocab[\"<PAD>\"]\n    ids = [vocab.get(t, unk) for t in tokens]\n    if len(ids) > max_len:\n        ids = ids[:max_len]\n    else:\n        ids += [pad] * (max_len - len(ids))\n    return ids\n\n\n# ============================================================\n# Load & split train samples\n# ============================================================\n\nstart_time = time.time()\n\nprint(f\"Loading up to {MAX_TRAIN_SAMPLES} training samples from {TRAIN_PATH}\")\nall_samples = list(read_jsonl(TRAIN_PATH, max_lines=MAX_TRAIN_SAMPLES))\nprint(f\"Loaded {len(all_samples)} training samples from {TRAIN_PATH}\")\n\nrandom.shuffle(all_samples)\n\nnum_all = len(all_samples)\nnum_val = int(num_all * VAL_FRACTION)\nnum_train = num_all - num_val\n\ntrain_samples = all_samples[:num_train]\nval_samples = all_samples[num_train:]\n\nprint(f\"Train samples: {len(train_samples)}, Validation samples: {len(val_samples)}\")\n\n\n# ============================================================\n# Build vocabulary\n# ============================================================\n\nvocab = {\"<PAD>\": 0, \"<UNK>\": 1}\nvocab_next_id = 2\n\n\ndef add_tokens_to_vocab(tokens):\n    global vocab_next_id\n    for t in tokens:\n        if t not in vocab:\n            vocab[t] = vocab_next_id\n            vocab_next_id += 1\n\n\nprint(\"Building vocabulary...\")\nfor sample in all_samples:\n    q_tokens = tokenize(sample[\"question_text\"])\n    add_tokens_to_vocab(q_tokens)\n\n    doc_tokens = sample[\"document_text\"].split()\n    for cand in sample[\"long_answer_candidates\"]:\n        st, end = cand[\"start_token\"], cand[\"end_token\"]\n        cand_tokens = doc_tokens[st:end]\n        add_tokens_to_vocab([t.lower() for t in cand_tokens])\n\nvocab_size = len(vocab)\nprint(f\"Vocab size: {vocab_size}\")\n\n\n# ============================================================\n# Build training pairs (question–candidate)\n# ============================================================\n\ndef build_training_pairs(samples, vocab, max_neg_per_q):\n    q_ids_list, p_ids_list, labels_list = [], [], []\n\n    for sample in samples:\n        doc_tokens = sample[\"document_text\"].split()\n        q_tokens = tokenize(sample[\"question_text\"])\n        q_ids_single = tokens_to_ids(q_tokens, vocab, MAX_Q_LEN)\n\n        annotations = sample.get(\"annotations\", [])\n        gold_indices = set()\n        if annotations:\n            ann = annotations[0]\n            long_answer = ann.get(\"long_answer\", {})\n            cand_idx = long_answer.get(\"candidate_index\", -1)\n            if cand_idx is not None and cand_idx >= 0:\n                gold_indices.add(int(cand_idx))\n\n        cands = sample[\"long_answer_candidates\"]\n        if not cands:\n            continue\n\n        pos_indices, neg_indices = [], []\n        for i, cand in enumerate(cands):\n            if i in gold_indices:\n                pos_indices.append(i)\n            else:\n                neg_indices.append(i)\n\n        kept_indices = list(pos_indices)\n        if neg_indices:\n            random.shuffle(neg_indices)\n            kept_indices.extend(neg_indices[:max_neg_per_q])\n\n        if not kept_indices:\n            continue\n\n        for i in kept_indices:\n            cand = cands[i]\n            st, end = cand[\"start_token\"], cand[\"end_token\"]\n            cand_tokens = [t.lower() for t in doc_tokens[st:end]]\n            p_ids = tokens_to_ids(cand_tokens, vocab, MAX_P_LEN)\n\n            label = 1 if i in gold_indices else 0\n\n            q_ids_list.append(q_ids_single)\n            p_ids_list.append(p_ids)\n            labels_list.append(label)\n\n    q_ids_tensor = torch.tensor(np.array(q_ids_list), dtype=torch.long)\n    p_ids_tensor = torch.tensor(np.array(p_ids_list), dtype=torch.long)\n    labels_tensor = torch.tensor(np.array(labels_list), dtype=torch.float32).unsqueeze(1)\n    return q_ids_tensor, p_ids_tensor, labels_tensor\n\n\nprint(\"Building training pairs...\")\ntrain_q_ids, train_p_ids, train_labels = build_training_pairs(\n    train_samples, vocab, MAX_NEG_PER_Q\n)\nprint(\n    f\"Prepared {len(train_labels)} training examples \"\n    f\"(question-paragraph candidate pairs)\"\n)\n\n\n# ============================================================\n# Model\n# ============================================================\n\nclass BiLSTMQA(nn.Module):\n    def __init__(self, vocab_size, embed_dim, hidden_dim):\n        super().__init__()\n        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n        self.q_lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)\n        self.p_lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)\n        self.fc = nn.Sequential(\n            nn.Linear(hidden_dim * 4, 128),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(128, 1),\n        )\n\n    def forward(self, q_ids, p_ids):\n        q_emb = self.embedding(q_ids)\n        p_emb = self.embedding(p_ids)\n\n        q_out, _ = self.q_lstm(q_emb)\n        p_out, _ = self.p_lstm(p_emb)\n\n        q_repr, _ = torch.max(q_out, dim=1)\n        p_repr, _ = torch.max(p_out, dim=1)\n\n        h = torch.cat([q_repr, p_repr], dim=1)\n        logits = self.fc(h)\n        return logits\n\n\nmodel = BiLSTMQA(vocab_size, EMBED_DIM, HIDDEN_DIM).to(device)\npos_weight = torch.tensor([3.0], device=device)\ncriterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\noptimizer = optim.Adam(model.parameters(), lr=LR)\n\n\n# ============================================================\n# Validation helper (NOTE: no model.eval() here)\n# ============================================================\n\ndef eval_long_answer_accuracy(model, samples, vocab):\n    correct = 0\n    total = 0\n\n    with torch.no_grad():\n        for sample in samples:\n            doc_tokens = sample[\"document_text\"].split()\n            q_tokens = tokenize(sample[\"question_text\"])\n            q_ids_single = tokens_to_ids(q_tokens, vocab, MAX_Q_LEN)\n\n            candidates = sample[\"long_answer_candidates\"]\n            if not candidates:\n                continue\n\n            cand_spans = []\n            cand_p_ids = []\n\n            for cand in candidates:\n                st, end = cand[\"start_token\"], cand[\"end_token\"]\n                cand_spans.append((st, end))\n                cand_tokens = [t.lower() for t in doc_tokens[st:end]]\n                cand_p_ids.append(tokens_to_ids(cand_tokens, vocab, MAX_P_LEN))\n\n            q_batch = torch.tensor(\n                np.repeat([q_ids_single], len(cand_spans), axis=0),\n                dtype=torch.long,\n                device=device,\n            )\n            p_batch = torch.tensor(np.array(cand_p_ids), dtype=torch.long, device=device)\n\n            logits = model(q_batch, p_batch).squeeze(1)\n            probs = torch.sigmoid(logits).cpu().numpy()\n\n            best_idx = int(probs.argmax())\n            pred_cand_idx = best_idx\n\n            annotations = sample.get(\"annotations\", [])\n            gold_idx = -1\n            if annotations:\n                ann = annotations[0]\n                la = ann.get(\"long_answer\", {})\n                gold_idx = la.get(\"candidate_index\", -1)\n\n            if pred_cand_idx == gold_idx:\n                correct += 1\n            total += 1\n\n    acc = correct / total if total > 0 else 0.0\n    return acc, correct, total\n\n\n# ============================================================\n# Train (model stays in train mode; validation uses no_grad only)\n# ============================================================\n\nnum_train_pairs = train_labels.size(0)\nprint(f\"Number of training pairs: {num_train_pairs}\")\n\nbest_val_acc = 0.0\nbest_state_dict = None\npatience = 2\nno_improve_epochs = 0\n\nfor epoch in range(EPOCHS):\n    model.train()\n\n    perm = torch.randperm(num_train_pairs)\n    epoch_loss = 0.0\n\n    for start_idx in range(0, num_train_pairs, BATCH_SIZE):\n        idx = perm[start_idx : start_idx + BATCH_SIZE]\n        q_batch = train_q_ids[idx].to(device)\n        p_batch = train_p_ids[idx].to(device)\n        y_batch = train_labels[idx].to(device)\n\n        optimizer.zero_grad()\n        logits = model(q_batch, p_batch)\n        loss = criterion(logits, y_batch)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n        optimizer.step()\n\n        epoch_loss += loss.item() * q_batch.size(0)\n\n    epoch_loss /= num_train_pairs\n    print(f\"Epoch {epoch + 1}/{EPOCHS} - Loss: {epoch_loss:.4f}\")\n\n    val_acc, val_correct, val_total = eval_long_answer_accuracy(model, val_samples, vocab)\n    print(\n        f\"Validation long-answer accuracy (no threshold): {val_acc:.4f} \"\n        f\"({val_correct}/{val_total})\"\n    )\n\n    if val_acc > best_val_acc:\n        best_val_acc = val_acc\n        best_state_dict = {k: v.cpu() for k, v in model.state_dict().items()}\n        no_improve_epochs = 0\n    else:\n        no_improve_epochs += 1\n        if no_improve_epochs >= patience:\n            print(\"Early stopping: no improvement on validation.\")\n            break\n\nif best_state_dict is not None:\n    model.load_state_dict(best_state_dict)\n    model.to(device)\n    print(f\"Restored best model with val acc = {best_val_acc:.4f}\")\n\n\n# ============================================================\n# Inference on test set (now we can safely use eval())\n# ============================================================\n\nmodel.eval()\n\nlong_preds = {}\nshort_preds = {}\n\nprint(f\"Running inference on test set: {TEST_PATH}\")\ntest_count = 0\n\nwith torch.no_grad():\n    for sample in read_jsonl(TEST_PATH):\n        example_id = str(sample[\"example_id\"])\n        base_id = example_id\n\n        doc_tokens = sample[\"document_text\"].split()\n        q_tokens = tokenize(sample[\"question_text\"])\n        q_ids_single = tokens_to_ids(q_tokens, vocab, MAX_Q_LEN)\n\n        cand_spans = []\n        cand_p_ids = []\n\n        for cand in sample[\"long_answer_candidates\"]:\n            st, end = cand[\"start_token\"], cand[\"end_token\"]\n            cand_spans.append((st, end))\n            cand_tokens = [t.lower() for t in doc_tokens[st:end]]\n            cand_p_ids.append(tokens_to_ids(cand_tokens, vocab, MAX_P_LEN))\n\n        if len(cand_spans) == 0:\n            long_preds[base_id] = \"\"\n            short_preds[base_id] = \"\"\n            test_count += 1\n            if test_count % 100 == 0:\n                print(f\"Processed {test_count} test samples...\")\n            continue\n\n        q_batch = torch.tensor(\n            np.repeat([q_ids_single], len(cand_spans), axis=0),\n            dtype=torch.long,\n            device=device,\n        )\n        p_batch = torch.tensor(np.array(cand_p_ids), dtype=torch.long, device=device)\n\n        logits = model(q_batch, p_batch).squeeze(1)\n        probs = torch.sigmoid(logits).cpu().numpy()\n\n        best_idx = int(probs.argmax())\n        best_start, best_end = cand_spans[best_idx]\n\n        long_pred_str = f\"{best_start}:{best_end}\"\n\n        long_text = \" \".join(doc_tokens[best_start:best_end]).lower()\n        if \" yes \" in (\" \" + long_text + \" \"):\n            short_pred_str = \"YES\"\n        elif \" no \" in (\" \" + long_text + \" \"):\n            short_pred_str = \"NO\"\n        else:\n            short_pred_str = \"\"\n\n        long_preds[base_id] = long_pred_str\n        short_preds[base_id] = short_pred_str\n\n        test_count += 1\n        if test_count % 100 == 0:\n            print(f\"Processed {test_count} test samples...\")\n\nprint(f\"Finished inference on {test_count} test samples.\")\n\n\n# ============================================================\n# Build submission\n# ============================================================\n\ndef base_from_row_id(row_id: str) -> str:\n    s = str(row_id)\n    if \"_\" in s:\n        return s.rsplit(\"_\", 1)[0]\n    return s\n\n\ndef build_submission_from_sample(sample_sub_path, long_preds, short_preds):\n    sub = pd.read_csv(sample_sub_path)\n    if \"PredictionString\" not in sub.columns or \"example_id\" not in sub.columns:\n        raise ValueError(\n            \"sample_submission.csv must have columns 'example_id' and 'PredictionString'.\"\n        )\n\n    sub[\"PredictionString\"] = \"\"\n\n    is_long = sub[\"example_id\"].astype(str).str.endswith(\"_long\")\n    is_short = sub[\"example_id\"].astype(str).str.endswith(\"_short\")\n\n    long_ids = sub.loc[is_long, \"example_id\"].astype(str)\n    sub.loc[is_long, \"PredictionString\"] = [\n        long_preds.get(base_from_row_id(eid), \"\") for eid in long_ids\n    ]\n\n    short_ids = sub.loc[is_short, \"example_id\"].astype(str)\n    sub.loc[is_short, \"PredictionString\"] = [\n        short_preds.get(base_from_row_id(eid), \"\") for eid in short_ids\n    ]\n\n    return sub\n\n\nprint(f\"Building submission from {SAMPLE_SUB_PATH}\")\nsubmission_df = build_submission_from_sample(SAMPLE_SUB_PATH, long_preds, short_preds)\nsubmission_df.to_csv(OUTPUT_PATH, index=False)\nprint(f\"Saved submission to {OUTPUT_PATH}\")\n\nelapsed = time.time() - start_time\nmins = int(elapsed // 60)\nsecs = int(elapsed % 60)\nprint(f\"Execution time: {mins} minutes {secs} seconds.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T17:55:49.223625Z","iopub.execute_input":"2025-12-13T17:55:49.223932Z","iopub.status.idle":"2025-12-13T18:33:22.308614Z","shell.execute_reply.started":"2025-12-13T17:55:49.223908Z","shell.execute_reply":"2025-12-13T18:33:22.307878Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}