{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"a49e33f2-b768-4d72-85aa-7017db65460a","cell_type":"code","source":"\nimport os, random, numpy as np, pandas as pd, torch, torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\n\nSEED = 42\nMAX_LEN = 2048\ndef seed_everything(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nseed_everything(SEED)\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Using device:\", device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T13:58:46.010464Z","iopub.execute_input":"2026-02-15T13:58:46.011231Z","iopub.status.idle":"2026-02-15T13:58:46.018463Z","shell.execute_reply.started":"2026-02-15T13:58:46.011194Z","shell.execute_reply":"2026-02-15T13:58:46.017852Z"}},"outputs":[],"execution_count":null},{"id":"5e9903e2-f762-494a-8acd-410d49d2b4f6","cell_type":"code","source":"\nROOT = \"/kaggle/input/competitions/stanford-rna-3d-folding-2\"\n\ntrain_seq = pd.read_csv(os.path.join(ROOT, \"train_sequences.csv\"))\ntrain_lab = pd.read_csv(os.path.join(ROOT, \"train_labels.csv\"), low_memory=False)\nval_seq   = pd.read_csv(os.path.join(ROOT, \"validation_sequences.csv\"))\nval_lab   = pd.read_csv(os.path.join(ROOT, \"validation_labels.csv\"), low_memory=False)\ntest_seq  = pd.read_csv(os.path.join(ROOT, \"test_sequences.csv\"))\n\ntrain_lab[\"target_id\"] = train_lab[\"ID\"].str.split(\"_\").str[0]\nval_lab[\"target_id\"]   = val_lab[\"ID\"].str.split(\"_\").str[0]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T14:03:03.953397Z","iopub.execute_input":"2026-02-15T14:03:03.954109Z","iopub.status.idle":"2026-02-15T14:03:24.039486Z","shell.execute_reply.started":"2026-02-15T14:03:03.954073Z","shell.execute_reply":"2026-02-15T14:03:24.038691Z"}},"outputs":[],"execution_count":null},{"id":"2fb72f32-c973-4992-8c5f-39a8343263ba","cell_type":"code","source":"\ndef build_samples(seq_df, lab_df):\n    seq_map = dict(zip(seq_df[\"target_id\"], seq_df[\"sequence\"]))\n    samples = []\n\n    for tid, g in lab_df.groupby(\"target_id\"):\n        if tid not in seq_map:\n            continue\n\n        full_seq = seq_map[tid]\n        n_copies = g[\"copy\"].max()\n        L_total = len(full_seq)\n\n        if L_total % n_copies != 0:\n            continue\n\n        L_copy = L_total // n_copies\n\n        for c in range(1, n_copies + 1):\n            g_c = g[g[\"copy\"] == c]\n\n            coords = np.zeros((L_copy, 3), dtype=np.float32)\n            mask   = np.zeros((L_copy,), dtype=np.float32)\n\n            r = g_c[\"resid\"].astype(int).to_numpy() - 1\n            xyz = g_c[[\"x_1\",\"y_1\",\"z_1\"]].to_numpy(np.float32)\n\n            ok = (r >= 0) & (r < L_copy)\n            r = r[ok]\n            xyz = xyz[ok]\n\n            coords[r] = xyz\n            mask[r] = 1.0\n\n            seq_slice = full_seq[(c-1)*L_copy : c*L_copy]\n\n            # LENGTH CAP\n            if len(seq_slice) > MAX_LEN:\n                seq_slice = seq_slice[:MAX_LEN]\n                coords = coords[:MAX_LEN]\n                mask = mask[:MAX_LEN]\n\n            if mask.sum() == 0:\n                continue\n            \n            samples.append((seq_slice, coords, mask))\n\n\n    return samples\n\ntrain_samples = build_samples(train_seq, train_lab)\nval_samples   = build_samples(val_seq, val_lab)\n\nprint(\"Train samples:\", len(train_samples))\nprint(\"Val samples  :\", len(val_samples))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T14:03:26.382075Z","iopub.execute_input":"2026-02-15T14:03:26.382821Z","iopub.status.idle":"2026-02-15T14:03:38.098024Z","shell.execute_reply.started":"2026-02-15T14:03:26.382792Z","shell.execute_reply":"2026-02-15T14:03:38.097311Z"}},"outputs":[],"execution_count":null},{"id":"220cc85c-f26a-4bf0-9a0c-341a3e18e654","cell_type":"code","source":"\nVOCAB = {c:i+1 for i,c in enumerate(\"ACGU\")}\n\nclass RNADataset(Dataset):\n    def __init__(self, samples):\n        self.samples = samples\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        seq, coords, mask = self.samples[idx]\n        x = torch.tensor([VOCAB.get(c,0) for c in seq], dtype=torch.long)\n        y = torch.tensor(coords, dtype=torch.float32)\n        m = torch.tensor(mask, dtype=torch.float32)\n        return x, y, m\n\ndef collate_fn(batch):\n    xs, ys, ms = zip(*batch)\n    max_len = max(x.size(0) for x in xs)\n\n    x_pad = torch.zeros(len(xs), max_len, dtype=torch.long)\n    y_pad = torch.zeros(len(xs), max_len, 3)\n    m_pad = torch.zeros(len(xs), max_len)\n\n    for i,(x,y,m) in enumerate(zip(xs,ys,ms)):\n        L = x.size(0)\n        x_pad[i,:L] = x\n        y_pad[i,:L] = y\n        m_pad[i,:L] = m\n\n    return x_pad.to(device), y_pad.to(device), m_pad.to(device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T14:04:44.737052Z","iopub.execute_input":"2026-02-15T14:04:44.738013Z","iopub.status.idle":"2026-02-15T14:04:44.745914Z","shell.execute_reply.started":"2026-02-15T14:04:44.737979Z","shell.execute_reply":"2026-02-15T14:04:44.745374Z"}},"outputs":[],"execution_count":null},{"id":"b9431fe7-f644-410a-8751-18985370f2b2","cell_type":"code","source":"\ntrain_loader = DataLoader(RNADataset(train_samples), batch_size=4, shuffle=True, collate_fn=collate_fn)\nval_loader   = DataLoader(RNADataset(val_samples), batch_size=4, shuffle=False, collate_fn=collate_fn)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\nscaler = GradScaler()\n\ndef masked_mse(pred, target, mask):\n    diff = (pred - target) ** 2\n    diff = diff.sum(dim=-1)\n\n    valid = mask > 0.5\n    if valid.sum() == 0:\n        return torch.tensor(0.0, device=pred.device)\n\n    return diff[valid].mean()\n\nfor epoch in range(5):\n    model.train()\n    for x,y,m in train_loader:\n        optimizer.zero_grad()\n        with autocast():\n            pred = model(x)\n            loss = masked_mse(pred, y, m)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n    model.eval()\n    val_loss = 0\n    with torch.no_grad():\n        for x,y,m in val_loader:\n            pred = model(x)\n            val_loss += masked_mse(pred,y,m).item()\n\n    print(f\"Epoch {epoch+1} | Val Loss: {val_loss/len(val_loader):.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T14:04:49.137578Z","iopub.execute_input":"2026-02-15T14:04:49.138151Z","iopub.status.idle":"2026-02-15T14:09:22.930516Z","shell.execute_reply.started":"2026-02-15T14:04:49.138122Z","shell.execute_reply":"2026-02-15T14:09:22.929562Z"}},"outputs":[],"execution_count":null},{"id":"fdfc9a52-bff9-4571-be75-077db22c2678","cell_type":"code","source":"\nclass SimpleTransformer(nn.Module):\n    def __init__(self, vocab_size=5, d_model=128):\n        super().__init__()\n        self.embed = nn.Embedding(vocab_size, d_model, padding_idx=0)\n        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=8, batch_first=True)\n        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=4)\n        self.head = nn.Linear(d_model, 3)\n\n    def forward(self, x):\n        x = self.embed(x)\n        x = self.encoder(x)\n        return self.head(x)\n\nmodel = SimpleTransformer().to(device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T13:34:27.392298Z","iopub.execute_input":"2026-02-15T13:34:27.393014Z","iopub.status.idle":"2026-02-15T13:34:27.833588Z","shell.execute_reply.started":"2026-02-15T13:34:27.392986Z","shell.execute_reply":"2026-02-15T13:34:27.833031Z"}},"outputs":[],"execution_count":null},{"id":"0b3041c0-ad6e-4545-8715-95fe79989a7f","cell_type":"code","source":"total = 0\nlabeled = 0\n\nfor _, coords, mask in train_samples:\n    total += mask.shape[0]\n    labeled += mask.sum()\n\nprint(\"Train labeled fraction:\", labeled / total)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T13:46:13.723868Z","iopub.execute_input":"2026-02-15T13:46:13.724607Z","iopub.status.idle":"2026-02-15T13:46:13.762435Z","shell.execute_reply.started":"2026-02-15T13:46:13.724576Z","shell.execute_reply":"2026-02-15T13:46:13.761778Z"}},"outputs":[],"execution_count":null},{"id":"4dd06584-5100-48a3-a145-f8ea1c2bed0a","cell_type":"code","source":"lengths = [len(s[0]) for s in train_samples]\nprint(\"Max length:\", max(lengths))\nprint(\"Mean length:\", np.mean(lengths))\nprint(\"95th percentile:\", np.percentile(lengths, 95))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-15T13:59:43.716143Z","iopub.execute_input":"2026-02-15T13:59:43.716824Z","iopub.status.idle":"2026-02-15T13:59:43.72477Z","shell.execute_reply.started":"2026-02-15T13:59:43.716795Z","shell.execute_reply":"2026-02-15T13:59:43.724005Z"}},"outputs":[],"execution_count":null}]}