{"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":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14483971,"sourceType":"datasetVersion","datasetId":9251134}],"dockerImageVersionId":31234,"isInternetEnabled":true,"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":"2026-01-26T06:15:30.772801Z","iopub.execute_input":"2026-01-26T06:15:30.773147Z","iopub.status.idle":"2026-01-26T06:15:31.168355Z","shell.execute_reply.started":"2026-01-26T06:15:30.773115Z","shell.execute_reply":"2026-01-26T06:15:31.167001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nfrom sklearn.preprocessing import normalize\nfrom tqdm import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:15:33.692150Z","iopub.execute_input":"2026-01-26T06:15:33.692662Z","iopub.status.idle":"2026-01-26T06:15:34.516793Z","shell.execute_reply.started":"2026-01-26T06:15:33.692627Z","shell.execute_reply":"2026-01-26T06:15:34.515402Z"}},"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-26T06:15:36.462075Z","iopub.execute_input":"2026-01-26T06:15:36.462605Z","iopub.status.idle":"2026-01-26T06:15:37.694567Z","shell.execute_reply.started":"2026-01-26T06:15:36.462570Z","shell.execute_reply":"2026-01-26T06:15:37.693384Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def clean_pid(pid):\n    # sp|A0A0C5B5G6|MOTSC_HUMAN → A0A0C5B5G6\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 embeddings:\", len(train_emb))\nprint(\"Clean test embeddings :\", len(test_emb))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:15:38.942137Z","iopub.execute_input":"2026-01-26T06:15:38.942503Z","iopub.status.idle":"2026-01-26T06:15:39.121525Z","shell.execute_reply.started":"2026-01-26T06:15:38.942473Z","shell.execute_reply":"2026-01-26T06:15:39.120255Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ids = list(train_emb.keys())\ntest_ids  = list(test_emb.keys())\n\nX_train = np.vstack([train_emb[p] for p in train_ids])\nX_test  = np.vstack([test_emb[p] for p in test_ids])\n\n# Normalize for cosine similarity\nX_train = normalize(X_train)\nX_test  = normalize(X_test)\n\nprint(\"X_train shape:\", X_train.shape)\nprint(\"X_test shape :\", X_test.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:15:41.592130Z","iopub.execute_input":"2026-01-26T06:15:41.592495Z","iopub.status.idle":"2026-01-26T06:15:42.454173Z","shell.execute_reply.started":"2026-01-26T06:15:41.592466Z","shell.execute_reply":"2026-01-26T06:15:42.453242Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nDATA_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-26T06:15:45.037076Z","iopub.execute_input":"2026-01-26T06:15:45.037462Z","iopub.status.idle":"2026-01-26T06:15:45.351190Z","shell.execute_reply.started":"2026-01-26T06:15:45.037431Z","shell.execute_reply":"2026-01-26T06:15:45.350404Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import defaultdict\n\nprot2go = defaultdict(list)\n\nfor _, row in train_terms.iterrows():\n    prot2go[row[\"protein_id\"]].append(row[\"go_id\"])\n\nprint(\"Proteins with labels:\", len(prot2go))\nprint(\"Example mapping:\", list(prot2go.items())[:1])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:15:48.192956Z","iopub.execute_input":"2026-01-26T06:15:48.193311Z","iopub.status.idle":"2026-01-26T06:16:11.064695Z","shell.execute_reply.started":"2026-01-26T06:15:48.193264Z","shell.execute_reply":"2026-01-26T06:16:11.063681Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics.pairwise import cosine_similarity\nfrom collections import defaultdict\n\nTOP_NEIGHBORS = 10    # start safe\nCHUNK = 300           # memory-safe\nOUTPUT_PATH = \"/kaggle/working/submission.tsv\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:16:57.014378Z","iopub.execute_input":"2026-01-26T06:16:57.015444Z","iopub.status.idle":"2026-01-26T06:16:57.045569Z","shell.execute_reply.started":"2026-01-26T06:16:57.015406Z","shell.execute_reply":"2026-01-26T06:16:57.044573Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install faiss-cpu\nimport faiss\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:19:01.357085Z","iopub.execute_input":"2026-01-26T06:19:01.357509Z","iopub.status.idle":"2026-01-26T06:19:05.295434Z","shell.execute_reply.started":"2026-01-26T06:19:01.357464Z","shell.execute_reply":"2026-01-26T06:19:05.294343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dim = X_train.shape[1]\n\nindex = faiss.IndexFlatIP(dim)   # inner product = cosine (normalized vectors)\nindex.add(X_train.astype(np.float32))\n\nprint(\"FAISS index size:\", index.ntotal)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:19:12.673192Z","iopub.execute_input":"2026-01-26T06:19:12.674542Z","iopub.status.idle":"2026-01-26T06:19:12.819268Z","shell.execute_reply.started":"2026-01-26T06:19:12.674497Z","shell.execute_reply":"2026-01-26T06:19:12.818426Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"faiss.__version__\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:19:15.068793Z","iopub.execute_input":"2026-01-26T06:19:15.069100Z","iopub.status.idle":"2026-01-26T06:19:15.075987Z","shell.execute_reply.started":"2026-01-26T06:19:15.069074Z","shell.execute_reply":"2026-01-26T06:19:15.075135Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"FAISS_K = 100     # retrieve top 100 neighbors\nUSE_TOP = 15      # actually aggregate only best 15\n\nD, I = index.search(X_test.astype(np.float32), FAISS_K)\n\nprint(\"Search done:\", D.shape, I.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T06:19:31.326506Z","iopub.execute_input":"2026-01-26T06:19:31.326961Z","iopub.status.idle":"2026-01-26T06:23:10.643118Z","shell.execute_reply.started":"2026-01-26T06:19:31.326915Z","shell.execute_reply":"2026-01-26T06:23:10.641798Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"OUT_PATH = \"/kaggle/working/submission.tsv\"\n\nwith open(OUT_PATH, \"w\") as fout:\n    for i in tqdm(range(len(test_ids))):\n        pid = test_ids[i]\n        idxs = I[i][:USE_TOP]\n        sims = D[i][:USE_TOP]\n\n        score_map = defaultdict(float)\n\n        for idx, sim in zip(idxs, sims):\n            train_pid = train_ids[idx]\n            weight = sim\n\n            for go in prot2go.get(train_pid, []):\n                score_map[go] += weight\n\n        total = sum(score_map.values()) + 1e-8\n        for go, sc in score_map.items():\n            fout.write(f\"{pid}\\t{go}\\t{round(sc/total,3)}\\n\")\n\nprint(\"Saved:\", OUT_PATH)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T07:02:59.257683Z","iopub.execute_input":"2026-01-26T07:02:59.258223Z","iopub.status.idle":"2026-01-26T07:05:14.297157Z","shell.execute_reply.started":"2026-01-26T07:02:59.258185Z","shell.execute_reply":"2026-01-26T07:05:14.294950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}