{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14483971,"sourceType":"datasetVersion","datasetId":9251134}],"dockerImageVersionId":31240,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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":"2026-01-15T08:37:50.114015Z","iopub.execute_input":"2026-01-15T08:37:50.114687Z","iopub.status.idle":"2026-01-15T08:37:50.134200Z","shell.execute_reply.started":"2026-01-15T08:37:50.114657Z","shell.execute_reply":"2026-01-15T08:37:50.133643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom sklearn.preprocessing import normalize, MultiLabelBinarizer\nfrom sklearn.model_selection import train_test_split\nfrom collections import defaultdict\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:37:50.135269Z","iopub.execute_input":"2026-01-15T08:37:50.135512Z","iopub.status.idle":"2026-01-15T08:37:53.571247Z","shell.execute_reply.started":"2026-01-15T08:37:50.135473Z","shell.execute_reply":"2026-01-15T08:37:53.570447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"EMB_PATH = \"/kaggle/input/embeddings\"\n\ntrain_emb_raw = np.load(\n    f\"{EMB_PATH}/esm2_train_embeddings.npy\",\n    allow_pickle=True\n).item()\n\ntest_emb_raw = np.load(\n    f\"{EMB_PATH}/esm2_test_embeddings.npy\",\n    allow_pickle=True\n).item()\n\nprint(\"Train embeddings:\", len(train_emb_raw))\nprint(\"Test embeddings :\", len(test_emb_raw))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:39:15.598654Z","iopub.execute_input":"2026-01-15T08:39:15.599201Z","iopub.status.idle":"2026-01-15T08:39:21.472852Z","shell.execute_reply.started":"2026-01-15T08:39:15.599176Z","shell.execute_reply":"2026-01-15T08:39:21.472059Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def clean_pid(pid):\n    if \"|\" in pid:\n        return pid.split(\"|\")[1]\n    return pid\n\ntrain_emb = {clean_pid(k): v for k, v in train_emb_raw.items()}\ntest_emb  = {clean_pid(k): v for k, v in test_emb_raw.items()}\n\nprint(\"Clean train:\", len(train_emb))\nprint(\"Clean test :\", len(test_emb))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:39:55.948476Z","iopub.execute_input":"2026-01-15T08:39:55.949042Z","iopub.status.idle":"2026-01-15T08:39:56.049891Z","shell.execute_reply.started":"2026-01-15T08:39:55.949017Z","shell.execute_reply":"2026-01-15T08:39:56.049111Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ids = list(train_emb.keys())\ntest_ids  = list(test_emb.keys())\n\nX_all_train = np.vstack([train_emb[p] for p in train_ids])\nX_test = np.vstack([test_emb[p] for p in test_ids])\n\nX_all_train = normalize(X_all_train)\nX_test = normalize(X_test)\n\nprint(\"Train matrix:\", X_all_train.shape)\nprint(\"Test matrix :\", X_test.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:39:58.508055Z","iopub.execute_input":"2026-01-15T08:39:58.508325Z","iopub.status.idle":"2026-01-15T08:39:59.261081Z","shell.execute_reply.started":"2026-01-15T08:39:58.508305Z","shell.execute_reply":"2026-01-15T08:39:59.258217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DATA_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train\"\n\ntrain_terms = pd.read_csv(\n    f\"{DATA_PATH}/train_terms.tsv\",\n    sep=\"\\t\"\n)\n\ntrain_terms = train_terms.rename(columns={\n    \"EntryID\": \"protein_id\",\n    \"term\": \"go_id\",\n    \"aspect\": \"ontology\"\n})\n\nprint(\"Total annotations:\", len(train_terms))\ntrain_terms.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:40:02.238206Z","iopub.execute_input":"2026-01-15T08:40:02.238472Z","iopub.status.idle":"2026-01-15T08:40:02.455359Z","shell.execute_reply.started":"2026-01-15T08:40:02.238454Z","shell.execute_reply":"2026-01-15T08:40:02.454651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TOP_LABELS = 300\n\ntop_go = (\n    train_terms[\"go_id\"]\n    .value_counts()\n    .head(TOP_LABELS)\n    .index\n    .tolist()\n)\n\nprint(\"Using GO labels:\", len(top_go))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:40:05.897857Z","iopub.execute_input":"2026-01-15T08:40:05.898159Z","iopub.status.idle":"2026-01-15T08:40:05.947793Z","shell.execute_reply.started":"2026-01-15T08:40:05.898128Z","shell.execute_reply":"2026-01-15T08:40:05.947157Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"filtered = train_terms[train_terms[\"go_id\"].isin(top_go)]\n\nprot2labels = (\n    filtered.groupby(\"protein_id\")[\"go_id\"]\n    .apply(list)\n    .to_dict()\n)\n\nX, Y = [], []\n\nfor pid in train_ids:\n    if pid in prot2labels:\n        X.append(train_emb[pid])\n        Y.append(prot2labels[pid])\n\nmlb = MultiLabelBinarizer(classes=top_go)\nY = mlb.fit_transform(Y)\nX = np.array(X)\n\nprint(\"Final dataset:\", X.shape, Y.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:40:07.908311Z","iopub.execute_input":"2026-01-15T08:40:07.908588Z","iopub.status.idle":"2026-01-15T08:40:09.375407Z","shell.execute_reply.started":"2026-01-15T08:40:07.908569Z","shell.execute_reply":"2026-01-15T08:40:09.374764Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_train, X_val, y_train, y_val = train_test_split(\n    X, Y, test_size=0.1, random_state=42\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:40:13.307982Z","iopub.execute_input":"2026-01-15T08:40:13.308576Z","iopub.status.idle":"2026-01-15T08:40:13.413170Z","shell.execute_reply.started":"2026-01-15T08:40:13.308553Z","shell.execute_reply":"2026-01-15T08:40:13.412610Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", device)\n\nclass ProteinNet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(320, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, TOP_LABELS)\n        )\n\n    def forward(self, x):\n        return self.net(x)\n\nmodel = ProteinNet().to(device)\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:40:15.752634Z","iopub.execute_input":"2026-01-15T08:40:15.753239Z","iopub.status.idle":"2026-01-15T08:40:18.325184Z","shell.execute_reply.started":"2026-01-15T08:40:15.753218Z","shell.execute_reply":"2026-01-15T08:40:18.324387Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BATCH = 256\n\ntrain_loader = DataLoader(\n    TensorDataset(\n        torch.tensor(X_train).float(),\n        torch.tensor(y_train).float()\n    ),\n    batch_size=BATCH,\n    shuffle=True\n)\n\nval_loader = DataLoader(\n    TensorDataset(\n        torch.tensor(X_val).float(),\n        torch.tensor(y_val).float()\n    ),\n    batch_size=BATCH\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:41:02.784115Z","iopub.execute_input":"2026-01-15T08:41:02.784971Z","iopub.status.idle":"2026-01-15T08:41:02.925900Z","shell.execute_reply.started":"2026-01-15T08:41:02.784949Z","shell.execute_reply":"2026-01-15T08:41:02.925229Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"EPOCHS = 5\n\nfor epoch in range(EPOCHS):\n    model.train()\n    total_loss = 0\n\n    for xb, yb in train_loader:\n        xb, yb = xb.to(device), yb.to(device)\n\n        optimizer.zero_grad()\n        preds = model(xb)\n        loss = criterion(preds, yb)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    print(f\"Epoch {epoch+1} | Train loss: {total_loss:.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:41:15.963980Z","iopub.execute_input":"2026-01-15T08:41:15.964695Z","iopub.status.idle":"2026-01-15T08:41:20.644032Z","shell.execute_reply.started":"2026-01-15T08:41:15.964666Z","shell.execute_reply":"2026-01-15T08:41:20.643368Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\npreds_all = []\n\nwith torch.no_grad():\n    for i in tqdm(range(0, len(X_test), BATCH)):\n        xb = torch.tensor(X_test[i:i+BATCH]).float().to(device)\n        logits = model(xb)\n        probs = torch.sigmoid(logits).cpu().numpy()\n        preds_all.append(probs)\n\npreds_all = np.vstack(preds_all)\nprint(\"Pred matrix:\", preds_all.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:41:31.829579Z","iopub.execute_input":"2026-01-15T08:41:31.830294Z","iopub.status.idle":"2026-01-15T08:41:32.457842Z","shell.execute_reply.started":"2026-01-15T08:41:31.830271Z","shell.execute_reply":"2026-01-15T08:41:32.457054Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TOP_K = 30\nOUT_PATH = \"/kaggle/working/submission.tsv\"\n\nwith open(OUT_PATH, \"w\") as f:\n    for i, pid in enumerate(test_ids):\n        top_idx = preds_all[i].argsort()[-TOP_K:]\n        for idx in top_idx:\n            score = preds_all[i][idx]\n            if score > 0.01:\n                f.write(f\"{pid}\\t{mlb.classes_[idx]}\\t{round(float(score),3)}\\n\")\n\nprint(\"Saved:\", OUT_PATH)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T08:42:17.677836Z","iopub.execute_input":"2026-01-15T08:42:17.678574Z","iopub.status.idle":"2026-01-15T08:42:43.076464Z","shell.execute_reply.started":"2026-01-15T08:42:17.678550Z","shell.execute_reply":"2026-01-15T08:42:43.075841Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}