{"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":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, re, json\nimport numpy as np\nimport pandas as pd\n\nDATASET_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\nTRAIN_DIR = os.path.join(DATASET_DIR, \"Train\")\nTEST_DIR  = os.path.join(DATASET_DIR, \"Test\")\n\nTRAIN_FASTA = os.path.join(TRAIN_DIR, \"train_sequences.fasta\")\nTRAIN_TERMS = os.path.join(TRAIN_DIR, \"train_terms.tsv\")\nTRAIN_TAXON = os.path.join(TRAIN_DIR, \"train_taxonomy.tsv\")\n\nTEST_FASTA  = os.path.join(TEST_DIR, \"testsuperset.fasta\")\nSAMPLE_SUB  = os.path.join(DATASET_DIR, \"sample_submission.tsv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:23:40.216056Z","iopub.execute_input":"2026-01-18T08:23:40.216235Z","iopub.status.idle":"2026-01-18T08:23:41.220859Z","shell.execute_reply.started":"2026-01-18T08:23:40.216216Z","shell.execute_reply":"2026-01-18T08:23:41.220284Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Helpers (FASTA + ID normalize)******","metadata":{}},{"cell_type":"code","source":"def read_fasta_to_df(path: str) -> pd.DataFrame:\n    ids, seqs = [], []\n    cur_id, cur_seq = None, []\n    with open(path, \"r\", encoding=\"utf-8\") as f:\n        for raw in f:\n            line = raw.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if cur_id is not None:\n                    ids.append(cur_id)\n                    seqs.append(\"\".join(cur_seq))\n                cur_id = line[1:].split()[0]\n                cur_seq = []\n            else:\n                cur_seq.append(line)\n    if cur_id is not None:\n        ids.append(cur_id)\n        seqs.append(\"\".join(cur_seq))\n    df = pd.DataFrame({\"protein_id\": ids, \"sequence\": seqs})\n    df[\"length\"] = df[\"sequence\"].str.len()\n    return df\n\ndef normalize_protein_id(x: str) -> str:\n    x = str(x).strip().split()[0]\n    # UniProt style: sp|ACC|NAME or tr|ACC|NAME\n    if \"|\" in x:\n        parts = x.split(\"|\")\n        if len(parts) >= 2:\n            x = parts[1]\n        else:\n            x = parts[-1]\n    # remove version suffix like ABC123.1\n    if \".\" in x:\n        left, right = x.rsplit(\".\", 1)\n        if right.isdigit():\n            x = left\n    return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:23:41.222554Z","iopub.execute_input":"2026-01-18T08:23:41.222961Z","iopub.status.idle":"2026-01-18T08:23:41.229871Z","shell.execute_reply.started":"2026-01-18T08:23:41.222937Z","shell.execute_reply":"2026-01-18T08:23:41.229174Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Load train/test + labels**","metadata":{}},{"cell_type":"code","source":"train_seq = read_fasta_to_df(TRAIN_FASTA)\ntest_seq  = read_fasta_to_df(TEST_FASTA)\n\ntrain_seq[\"protein_id\"] = train_seq[\"protein_id\"].map(normalize_protein_id)\ntest_seq[\"protein_id\"]  = test_seq[\"protein_id\"].map(normalize_protein_id)\n\ntrain_terms = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\ntrain_terms_std = train_terms.rename(columns={\"EntryID\":\"protein_id\",\"term\":\"go_id\",\"aspect\":\"aspect\"}).copy()\ntrain_terms_std[\"protein_id\"] = train_terms_std[\"protein_id\"].map(normalize_protein_id)\n\ntrain = train_seq.copy()\ntest  = test_seq.copy()\n\nprint(\"train:\", train.shape, \"| test:\", test.shape, \"| terms:\", train_terms_std.shape)\nprint(\"train_terms_std cols:\", train_terms_std.columns.tolist())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:23:41.230822Z","iopub.execute_input":"2026-01-18T08:23:41.231167Z","iopub.status.idle":"2026-01-18T08:23:45.889894Z","shell.execute_reply.started":"2026-01-18T08:23:41.231138Z","shell.execute_reply":"2026-01-18T08:23:45.889153Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with open(SAMPLE_SUB, \"r\", encoding=\"utf-8\") as f:\n    for _ in range(3):\n        print(\"sample protein_id:\", f.readline().strip().split(\"\\t\")[0])\nprint(\"test protein_id:\", test[\"protein_id\"].iloc[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:23:45.890871Z","iopub.execute_input":"2026-01-18T08:23:45.891159Z","iopub.status.idle":"2026-01-18T08:23:45.900452Z","shell.execute_reply.started":"2026-01-18T08:23:45.891130Z","shell.execute_reply":"2026-01-18T08:23:45.899897Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Create train/val split with duplicate-sequence safety**","metadata":{}},{"cell_type":"code","source":"SEED = 42\nVAL_FRAC = 0.10\n\ntrain_ids_path = \"/kaggle/working/train_ids.txt\"\nval_ids_path   = \"/kaggle/working/val_ids.txt\"\n\nif os.path.exists(train_ids_path) and os.path.exists(val_ids_path):\n    train_ids = pd.read_csv(train_ids_path, header=None)[0].astype(str).tolist()\n    val_ids   = pd.read_csv(val_ids_path,   header=None)[0].astype(str).tolist()\n    print(\"Loaded existing split:\", len(train_ids), len(val_ids))\nelse:\n    seq_to_rep = train.groupby(\"sequence\")[\"protein_id\"].min().rename(\"rep_id\").reset_index()\n    train_rep  = train.merge(seq_to_rep, on=\"sequence\", how=\"left\")\n\n    rep_ids = train_rep[\"rep_id\"].drop_duplicates().values\n    rng = np.random.default_rng(SEED)\n    rng.shuffle(rep_ids)\n\n    n_val = int(len(rep_ids) * VAL_FRAC)\n    val_rep_ids = set(rep_ids[:n_val])\n    train_rep_ids = set(rep_ids[n_val:])\n\n    val_ids   = train_rep.loc[train_rep[\"rep_id\"].isin(val_rep_ids), \"protein_id\"].unique().tolist()\n    train_ids = train_rep.loc[train_rep[\"rep_id\"].isin(train_rep_ids), \"protein_id\"].unique().tolist()\n\n    pd.Series(train_ids).to_csv(train_ids_path, index=False, header=False)\n    pd.Series(val_ids).to_csv(val_ids_path, index=False, header=False)\n\n    print(\"Created split:\", len(train_ids), len(val_ids))\n    print(\"protein overlap:\", len(set(train_ids) & set(val_ids)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:23:45.901398Z","iopub.execute_input":"2026-01-18T08:23:45.901750Z","iopub.status.idle":"2026-01-18T08:23:49.721728Z","shell.execute_reply.started":"2026-01-18T08:23:45.901722Z","shell.execute_reply":"2026-01-18T08:23:49.721016Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Build GO label vocab + pid→label list maps**","metadata":{}},{"cell_type":"code","source":"train_terms_train = train_terms_std[train_terms_std[\"protein_id\"].isin(train_ids)][[\"protein_id\",\"go_id\"]].copy()\ngo_vocab = sorted(train_terms_train[\"go_id\"].unique())\ngo2i = {go:i for i,go in enumerate(go_vocab)}\ni2go = {i:go for go,i in go2i.items()}\nn_labels = len(go_vocab)\nprint(\"n_labels:\", n_labels)\n\ndef make_pid2labels(pids):\n    df = train_terms_std[train_terms_std[\"protein_id\"].isin(pids)][[\"protein_id\",\"go_id\"]].copy()\n    df = df[df[\"go_id\"].isin(go2i)]\n    df[\"y\"] = df[\"go_id\"].map(go2i).astype(int)\n    return df.groupby(\"protein_id\")[\"y\"].apply(list).to_dict()\n\ntrain_pid2labels = make_pid2labels(set(train_ids))\nval_pid2labels   = make_pid2labels(set(val_ids))\n\nprint(\"train proteins w/labels:\", len(train_pid2labels))\nprint(\"val proteins w/labels:\", len(val_pid2labels))\n\npath = \"/kaggle/working/go_vocab.json\"\nwith open(path, \"w\", encoding=\"utf-8\") as f:\n    json.dump(go_vocab, f)\n\nprint(\"saved to:\", path)\nprint(\"exists:\", os.path.exists(path))\nprint(\"size bytes:\", os.path.getsize(path))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:23:49.722649Z","iopub.execute_input":"2026-01-18T08:23:49.722962Z","iopub.status.idle":"2026-01-18T08:23:51.474782Z","shell.execute_reply.started":"2026-01-18T08:23:49.722938Z","shell.execute_reply":"2026-01-18T08:23:51.474022Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Train CNN**","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast, GradScaler\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"device:\", device)\n\nMAX_LEN = 512\nBATCH_SIZE = 128\nNUM_WORKERS = 2\nEPOCHS = 3\nLR = 1e-3\n\nPAD = 0\nUNK = 1\n\nlut = np.full(256, UNK, dtype=np.uint8)\nalphabet = \"ACDEFGHIKLMNPQRSTVWY\"\nfor i, aa in enumerate(alphabet, start=2):\n    lut[ord(aa)] = i\n\ndef encode_fixed(seq: str) -> np.ndarray:\n    b = np.frombuffer(seq.encode(\"ascii\", \"ignore\"), dtype=np.uint8)\n    x = lut[b]\n    out = np.full((MAX_LEN,), PAD, dtype=np.uint8)\n    n = min(len(x), MAX_LEN)\n    out[:n] = x[:n]\n    return out\n\nclass ProteinDataset(Dataset):\n    def __init__(self, df, pid2labels):\n        self.pids = df[\"protein_id\"].astype(str).tolist()\n        self.pid2labels = pid2labels\n        self.X = np.stack([encode_fixed(s) for s in df[\"sequence\"].tolist()], axis=0)\n\n    def __len__(self):\n        return len(self.pids)\n\n    def __getitem__(self, idx):\n        pid = self.pids[idx]\n        x = self.X[idx]\n        labs = self.pid2labels.get(pid, [])\n        return x, labs\n\ndef collate_batch(batch):\n    xs, labs_list = zip(*batch)\n    X = torch.from_numpy(np.stack(xs, axis=0)).long()\n    Y = torch.zeros((len(xs), n_labels), dtype=torch.float32)\n    for i, labs in enumerate(labs_list):\n        if len(labs) > 0:\n            Y[i, labs] = 1.0\n    return X, Y\n\ntrain_df = train.loc[train[\"protein_id\"].isin(train_ids), [\"protein_id\",\"sequence\"]].copy()\nval_df   = train.loc[train[\"protein_id\"].isin(val_ids),   [\"protein_id\",\"sequence\"]].copy()\n\ntrain_loader = DataLoader(\n    ProteinDataset(train_df, train_pid2labels),\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    num_workers=NUM_WORKERS,\n    pin_memory=True,\n    collate_fn=collate_batch\n)\n\nclass CNNClassifier(nn.Module):\n    def __init__(self, vocab_size=256, emb_dim=128, ch=256, n_labels=0):\n        super().__init__()\n        self.emb = nn.Embedding(vocab_size, emb_dim, padding_idx=PAD)\n        self.conv3 = nn.Conv1d(emb_dim, ch, kernel_size=3, padding=1)\n        self.conv5 = nn.Conv1d(emb_dim, ch, kernel_size=5, padding=2)\n        self.conv7 = nn.Conv1d(emb_dim, ch, kernel_size=7, padding=3)\n        self.act = nn.ReLU()\n        self.pool = nn.AdaptiveMaxPool1d(1)\n        self.fc = nn.Linear(ch * 3, n_labels)\n\n    def forward(self, x):\n        e = self.emb(x).transpose(1, 2)\n        h3 = self.pool(self.act(self.conv3(e))).squeeze(-1)\n        h5 = self.pool(self.act(self.conv5(e))).squeeze(-1)\n        h7 = self.pool(self.act(self.conv7(e))).squeeze(-1)\n        return self.fc(torch.cat([h3, h5, h7], dim=1))\n\nmodel = CNNClassifier(n_labels=n_labels).to(device)\n\nfreq = train_terms_std[train_terms_std[\"protein_id\"].isin(train_ids)][\"go_id\"].value_counts()\npos = np.array([freq.get(go, 0) for go in go_vocab], dtype=np.float32)\nN = float(len(train_ids))\n\npos_weight = (N - pos) / (pos + 1e-6)\npos_weight = np.clip(pos_weight, 1.0, 50.0)  # cap is important\npos_weight_t = torch.tensor(pos_weight, dtype=torch.float32, device=device)\n\ncriterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight_t)\noptimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)\nscaler = GradScaler(\"cuda\", enabled=(device.type == \"cuda\"))\n\nprint(\"pos_weight min/mean/max:\", float(pos_weight_t.min()), float(pos_weight_t.mean()), float(pos_weight_t.max()))\n\n\nmodel.train()\nfor ep in range(1, EPOCHS + 1):\n    running = 0.0\n    for step, (X, Y) in enumerate(train_loader, start=1):\n        X = X.to(device, non_blocking=True)\n        Y = Y.to(device, non_blocking=True)\n\n        optimizer.zero_grad(set_to_none=True)\n\n        with autocast(\"cuda\", enabled=(device.type == \"cuda\")):\n            logits = model(X)\n            loss = criterion(logits, Y)\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n        running += float(loss.item())\n        if step % 200 == 0:\n            print(f\"epoch {ep} step {step}/{len(train_loader)} loss {running/200:.4f}\")\n            running = 0.0\n\nMODEL_PATH = \"/kaggle/working/cnn_go_model_balanced.pt\"\ntorch.save(model.state_dict(), MODEL_PATH)\nprint(\"Saved model:\", MODEL_PATH)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:36:27.147462Z","iopub.execute_input":"2026-01-18T08:36:27.147810Z","iopub.status.idle":"2026-01-18T08:37:40.157591Z","shell.execute_reply.started":"2026-01-18T08:36:27.147780Z","shell.execute_reply":"2026-01-18T08:37:40.156633Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast\n\nsmall_test = test.head(5000).copy()\n\nclass SmallTestDataset(Dataset):\n    def __init__(self, df):\n        self.pids = df[\"protein_id\"].astype(str).tolist()\n        self.X = np.stack([encode_fixed(s) for s in df[\"sequence\"].tolist()], axis=0)\n\n    def __len__(self): \n        return len(self.pids)\n\n    def __getitem__(self, idx):\n        return self.pids[idx], self.X[idx]\n\ndef collate_small(batch):\n    pids, xs = zip(*batch)\n    X = torch.from_numpy(np.stack(xs, axis=0)).long()\n    return list(pids), X\n\nsmall_loader = DataLoader(\n    SmallTestDataset(small_test),\n    batch_size=256,\n    shuffle=False,\n    num_workers=2,\n    pin_memory=True,\n    collate_fn=collate_small\n)\n\nTOPK = 200\nall_go = []\n\nmodel.eval()\nwith torch.no_grad():\n    for pids, X in small_loader:\n        X = X.to(device, non_blocking=True)\n        with autocast(\"cuda\", enabled=(device.type==\"cuda\")):\n            prob = torch.sigmoid(model(X))\n\n        _, idx = torch.topk(prob, k=TOPK, dim=1)\n        idx = idx.cpu().numpy()\n\n        for i in range(idx.shape[0]):\n            all_go.extend([i2go[int(j)] for j in idx[i]])\n\nprint(\"unique GO in sample:\", len(set(all_go)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:38:08.571299Z","iopub.execute_input":"2026-01-18T08:38:08.571915Z","iopub.status.idle":"2026-01-18T08:38:09.563583Z","shell.execute_reply.started":"2026-01-18T08:38:08.571879Z","shell.execute_reply":"2026-01-18T08:38:09.562885Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Trying to improve score","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\nclass TestDataset(Dataset):\n    def __init__(self, df):\n        self.pids = df[\"protein_id\"].astype(str).tolist()\n        self.X = np.stack([encode_fixed(s) for s in df[\"sequence\"].tolist()], axis=0)\n\n    def __len__(self):\n        return len(self.pids)\n\n    def __getitem__(self, idx):\n        return self.pids[idx], self.X[idx]\n\ndef collate_test(batch):\n    pids, xs = zip(*batch)\n    X = torch.from_numpy(np.stack(xs, axis=0)).long()\n    return list(pids), X\n\ntest_loader = DataLoader(\n    TestDataset(test[[\"protein_id\",\"sequence\"]].copy()),\n    batch_size=256,\n    shuffle=False,\n    num_workers=NUM_WORKERS,\n    pin_memory=True,\n    collate_fn=collate_test\n)\n\nprint(\"test_loader ready. batches:\", len(test_loader), \"| proteins:\", len(test))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:42:05.205503Z","iopub.execute_input":"2026-01-18T08:42:05.205836Z","iopub.status.idle":"2026-01-18T08:42:06.930804Z","shell.execute_reply.started":"2026-01-18T08:42:05.205812Z","shell.execute_reply":"2026-01-18T08:42:06.930060Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TOPK_SUB = 300\nTHR = 0.05\nOUT_PATH = \"/kaggle/working/submission_300_thr005.tsv\"\n\nmodel.eval()\nwritten = 0\n\nwith open(OUT_PATH, \"w\", encoding=\"utf-8\") as f:\n    with torch.no_grad():\n        for pids, X in test_loader:\n            X = X.to(device, non_blocking=True)\n            with autocast(\"cuda\", enabled=(device.type==\"cuda\")):\n                prob = torch.sigmoid(model(X))\n\n            sc, idx = torch.topk(prob, k=TOPK_SUB, dim=1)\n            sc = sc.cpu().numpy()\n            idx = idx.cpu().numpy()\n\n            for i, pid in enumerate(pids):\n                mask = sc[i] >= THR\n                js = np.where(mask)[0]\n                for j in js:\n                    go = i2go[int(idx[i, j])]\n                    s = float(sc[i, j])\n                    f.write(f\"{pid}\\t{go}\\t{s:.6f}\\n\")\n                    written += 1\n\nprint(\"Wrote:\", OUT_PATH, \"| rows:\", written)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:42:13.743643Z","iopub.execute_input":"2026-01-18T08:42:13.743985Z","iopub.status.idle":"2026-01-18T08:44:27.273686Z","shell.execute_reply.started":"2026-01-18T08:42:13.743962Z","shell.execute_reply":"2026-01-18T08:44:27.272720Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Build GO parent graph from Build GO parent graph from**","metadata":{}},{"cell_type":"code","source":"import re\nfrom collections import defaultdict, deque\n\nOBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n\nID_RE = re.compile(r\"^id:\\s+(GO:\\d{7})\")\nIS_A_RE = re.compile(r\"^is_a:\\s+(GO:\\d{7})\")\n\nparents = defaultdict(list)\nall_terms = set()\n\ncur = None\nwith open(OBO_PATH, \"r\", encoding=\"utf-8\") as f:\n    for line in f:\n        m = ID_RE.match(line)\n        if m:\n            cur = m.group(1)\n            all_terms.add(cur)\n            continue\n        m = IS_A_RE.match(line)\n        if m and cur:\n            parents[cur].append(m.group(1))\n\nprint(\"terms in obo:\", len(all_terms), \"| edges:\", sum(len(v) for v in parents.values()))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:46:47.923545Z","iopub.execute_input":"2026-01-18T08:46:47.924240Z","iopub.status.idle":"2026-01-18T08:46:48.470836Z","shell.execute_reply.started":"2026-01-18T08:46:47.924204Z","shell.execute_reply":"2026-01-18T08:46:48.470200Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Compute ancestors (cached) and propagate scores**","metadata":{}},{"cell_type":"code","source":"import numpy as np\n\nanc_cache = {}\n\ndef get_ancestors(term):\n    if term in anc_cache:\n        return anc_cache[term]\n    seen = set()\n    stack = parents.get(term, [])\n    while stack:\n        p = stack.pop()\n        if p in seen:\n            continue\n        seen.add(p)\n        stack.extend(parents.get(p, []))\n    anc_cache[term] = list(seen)\n    return anc_cache[term]\n\ndef propagate_scores(go_ids, scores):\n    \"\"\"\n    go_ids: list of GO terms predicted for one protein\n    scores: list of their scores\n    returns dict {go_term: score} with ancestor propagation via max\n    \"\"\"\n    out = {}\n    for go, s in zip(go_ids, scores):\n        if go not in out or s > out[go]:\n            out[go] = s\n        for anc in get_ancestors(go):\n            if anc not in out or s > out[anc]:\n                out[anc] = s\n    return out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:46:56.674539Z","iopub.execute_input":"2026-01-18T08:46:56.675061Z","iopub.status.idle":"2026-01-18T08:46:56.681198Z","shell.execute_reply.started":"2026-01-18T08:46:56.675035Z","shell.execute_reply":"2026-01-18T08:46:56.680535Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Predict test + write submission.tsv**","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast\nimport numpy as np\nimport torch\n\nSUB_PATH = \"/kaggle/working/submission.tsv\"\nTOPK_SUB = 300\nTHR = 0.05\nBATCH_SIZE = 256\nNUM_WORKERS = 2\n\nclass TestDataset(Dataset):\n    def __init__(self, df):\n        self.pids = df[\"protein_id\"].astype(str).tolist()\n        self.X = np.stack([encode_fixed(s) for s in df[\"sequence\"].tolist()], axis=0)\n\n    def __len__(self): return len(self.pids)\n    def __getitem__(self, idx): return self.pids[idx], self.X[idx]\n\ndef collate_test(batch):\n    pids, xs = zip(*batch)\n    X = torch.from_numpy(np.stack(xs, axis=0)).long()\n    return list(pids), X\n\ntest_loader = DataLoader(TestDataset(test[[\"protein_id\",\"sequence\"]].copy()),\n                         batch_size=BATCH_SIZE, shuffle=False,\n                         num_workers=NUM_WORKERS, pin_memory=True, collate_fn=collate_test)\n\nmodel.eval()\nwritten = 0\nwith open(SUB_PATH, \"w\", encoding=\"utf-8\") as f:\n    with torch.no_grad():\n        for pids, X in test_loader:\n            X = X.to(device, non_blocking=True)\n            with autocast(\"cuda\", enabled=(device.type==\"cuda\")):\n                prob = torch.sigmoid(model(X))\n\n            sc, idx = torch.topk(prob, k=TOPK_SUB, dim=1)\n            sc = sc.cpu().numpy()\n            idx = idx.cpu().numpy()\n\n            for i, pid in enumerate(pids):\n                go_ids = [i2go[int(j)] for j in idx[i]]\n                scores = sc[i].tolist()\n\n                prop = propagate_scores(go_ids, scores)\n\n                items = [(go, s) for go, s in prop.items() if s >= THR]\n                if not items:\n                    continue\n                items.sort(key=lambda x: x[1], reverse=True)\n                items = items[:TOPK_SUB]\n\n                for go, s in items:\n                    f.write(f\"{pid}\\t{go}\\t{s:.6f}\\n\")\n                    written += 1\n\n\nprint(\"Wrote:\", SUB_PATH, \"| rows:\", written)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:47:03.698562Z","iopub.execute_input":"2026-01-18T08:47:03.698853Z","iopub.status.idle":"2026-01-18T08:50:13.477475Z","shell.execute_reply.started":"2026-01-18T08:47:03.698831Z","shell.execute_reply":"2026-01-18T08:50:13.476454Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Validate format<!--  -->**","metadata":{}},{"cell_type":"code","source":"GO_RE = re.compile(r\"^GO:\\d{7}$\")\n\ndef validate_go_only(path, max_lines_to_check=200000):\n    bad = []\n    n = 0\n    with open(path, \"r\", encoding=\"utf-8\") as f:\n        for line in f:\n            n += 1\n            if max_lines_to_check and n > max_lines_to_check:\n                break\n            parts = line.rstrip(\"\\n\").split(\"\\t\")\n            if len(parts) != 3:\n                bad.append((n, f\"wrong_num_cols:{len(parts)}\"))\n                continue\n            pid, go, score = parts\n            if not GO_RE.match(go):\n                bad.append((n, f\"bad_go_id:{go}\"))\n                continue\n            try:\n                s = float(score)\n                if not (0.0 < s <= 1.0):\n                    bad.append((n, f\"score_out_of_range:{s}\"))\n            except Exception:\n                bad.append((n, f\"score_not_float:{score}\"))\n    return n, bad\n\nn_lines, problems = validate_go_only(SUB_PATH)\nprint(\"checked:\", n_lines)\nprint(\"problems:\", problems[:20])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:52:32.507451Z","iopub.execute_input":"2026-01-18T08:52:32.508061Z","iopub.status.idle":"2026-01-18T08:52:32.680490Z","shell.execute_reply.started":"2026-01-18T08:52:32.508032Z","shell.execute_reply":"2026-01-18T08:52:32.679967Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, pandas as pd\n\nSUB_PATH = \"/kaggle/working/submission_300_thr005.tsv\"\n\nprint(\"size MB:\", os.path.getsize(SUB_PATH)/(1024**2))\n\ndf = pd.read_csv(SUB_PATH, sep=\"\\t\", header=None, names=[\"protein_id\",\"go_id\",\"score\"])\nprint(\"rows:\", len(df))\nprint(\"unique proteins:\", df[\"protein_id\"].nunique(), \" / expected:\", len(test))\nprint(\"unique GO:\", df[\"go_id\"].nunique())\nprint(\"avg rows/protein:\", df.groupby(\"protein_id\").size().mean())\nprint(\"score min/mean/max:\", df[\"score\"].min(), df[\"score\"].mean(), df[\"score\"].max())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T08:52:43.592174Z","iopub.execute_input":"2026-01-18T08:52:43.592876Z","iopub.status.idle":"2026-01-18T08:53:17.317429Z","shell.execute_reply.started":"2026-01-18T08:52:43.592849Z","shell.execute_reply":"2026-01-18T08:53:17.316775Z"}},"outputs":[],"execution_count":null}]}