{"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,"isSourceIdPinned":false,"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-13T12:54:31.066616Z","iopub.execute_input":"2026-01-13T12:54:31.066933Z","iopub.status.idle":"2026-01-13T12:54:31.080225Z","shell.execute_reply.started":"2026-01-13T12:54:31.066899Z","shell.execute_reply":"2026-01-13T12:54:31.079278Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")\n\nfrom sklearn.exceptions import ConvergenceWarning\nwarnings.filterwarnings(\"ignore\", category=ConvergenceWarning)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:31.082114Z","iopub.execute_input":"2026-01-13T12:54:31.082448Z","iopub.status.idle":"2026-01-13T12:54:31.093150Z","shell.execute_reply.started":"2026-01-13T12:54:31.082415Z","shell.execute_reply":"2026-01-13T12:54:31.092295Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport os\n\nprint(os.listdir(\"/kaggle/input\"))\n\ntrain_emb = np.load(\n    \"/kaggle/input/embeddings/esm2_train_embeddings.npy\",\n    allow_pickle=True\n).item()\n\ntest_emb = np.load(\n    \"/kaggle/input/embeddings/esm2_test_embeddings.npy\",\n    allow_pickle=True\n).item()\n\nprint(len(train_emb), len(test_emb))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:31.094237Z","iopub.execute_input":"2026-01-13T12:54:31.094478Z","iopub.status.idle":"2026-01-13T12:54:31.942250Z","shell.execute_reply.started":"2026-01-13T12:54:31.094456Z","shell.execute_reply":"2026-01-13T12:54:31.941083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.linear_model import SGDClassifier\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom tqdm import tqdm\nimport os\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:31.944093Z","iopub.execute_input":"2026-01-13T12:54:31.944341Z","iopub.status.idle":"2026-01-13T12:54:31.950394Z","shell.execute_reply.started":"2026-01-13T12:54:31.944320Z","shell.execute_reply":"2026-01-13T12:54:31.949383Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# change dataset name if needed\nEMB_PATH = \"/kaggle/input/embeddings\"\n\ntrain_emb = np.load(f\"{EMB_PATH}/esm2_train_embeddings.npy\", allow_pickle=True).item()\ntest_emb  = np.load(f\"{EMB_PATH}/esm2_test_embeddings.npy\", allow_pickle=True).item()\n\nprint(\"Train embeddings:\", len(train_emb))\nprint(\"Test embeddings :\", len(test_emb))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:31.952006Z","iopub.execute_input":"2026-01-13T12:54:31.952388Z","iopub.status.idle":"2026-01-13T12:54:32.851810Z","shell.execute_reply.started":"2026-01-13T12:54:31.952350Z","shell.execute_reply":"2026-01-13T12:54:32.850892Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DATA_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train\"\n\ntrain_terms = pd.read_csv(f\"{DATA_PATH}/train_terms.tsv\", sep=\"\\t\")\ntrain_terms = train_terms.rename(columns={\n    \"EntryID\": \"protein_id\",\n    \"term\": \"go_id\",\n    \"aspect\": \"ontology\"\n})\n\ntrain_terms[\"ontology\"].value_counts()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:32.852653Z","iopub.execute_input":"2026-01-13T12:54:32.852915Z","iopub.status.idle":"2026-01-13T12:54:33.071566Z","shell.execute_reply.started":"2026-01-13T12:54:32.852896Z","shell.execute_reply":"2026-01-13T12:54:33.070826Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TOP_TERMS = 600  # per ontology (safe)\n\ntop_go_terms = {}\nfor ont in [\"F\", \"P\", \"C\"]:\n    top_go_terms[ont] = (\n        train_terms[train_terms[\"ontology\"] == ont][\"go_id\"]\n        .value_counts()\n        .head(TOP_TERMS)\n        .index\n        .tolist()\n    )\n\n{ont: len(v) for ont, v in top_go_terms.items()}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:33.072380Z","iopub.execute_input":"2026-01-13T12:54:33.072568Z","iopub.status.idle":"2026-01-13T12:54:33.273112Z","shell.execute_reply.started":"2026-01-13T12:54:33.072551Z","shell.execute_reply":"2026-01-13T12:54:33.271915Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"reduced_data = {}\n\nfor ont in [\"F\", \"P\", \"C\"]:\n    subset = train_terms[\n        (train_terms[\"ontology\"] == ont) &\n        (train_terms[\"go_id\"].isin(top_go_terms[ont]))\n    ]\n\n    prot2labels = (\n        subset.groupby(\"protein_id\")[\"go_id\"]\n        .apply(set)\n        .to_dict()\n    )\n\n    X, Y = [], []\n\n    for pid, emb in train_emb.items():\n        if pid in prot2labels:\n            X.append(emb)\n            Y.append(prot2labels[pid])\n\n    mlb = MultiLabelBinarizer(classes=top_go_terms[ont])\n    Y_bin = mlb.fit_transform(Y)\n\n    reduced_data[ont] = {\n        \"X\": np.array(X),\n        \"Y\": Y_bin,\n        \"mlb\": mlb\n    }\n\n    print(f\"{ont}: X={len(X)}, labels={Y_bin.shape[1]}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:33.274197Z","iopub.execute_input":"2026-01-13T12:54:33.274469Z","iopub.status.idle":"2026-01-13T12:54:35.795476Z","shell.execute_reply.started":"2026-01-13T12:54:33.274448Z","shell.execute_reply":"2026-01-13T12:54:35.794179Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def clean_protein_id(pid):\n    # UniProt FASTA format: sp|ACCESSION|NAME\n    if \"|\" in pid:\n        return pid.split(\"|\")[1]\n    return pid\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:35.796521Z","iopub.execute_input":"2026-01-13T12:54:35.796801Z","iopub.status.idle":"2026-01-13T12:54:35.801704Z","shell.execute_reply.started":"2026-01-13T12:54:35.796757Z","shell.execute_reply":"2026-01-13T12:54:35.800696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_emb_clean = {\n    clean_protein_id(pid): emb\n    for pid, emb in train_emb.items()\n}\n\ntest_emb_clean = {\n    clean_protein_id(pid): emb\n    for pid, emb in test_emb.items()\n}\n\nprint(\"Clean train embeddings:\", len(train_emb_clean))\nprint(\"Clean test embeddings :\", len(test_emb_clean))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:35.803935Z","iopub.execute_input":"2026-01-13T12:54:35.804195Z","iopub.status.idle":"2026-01-13T12:54:35.984040Z","shell.execute_reply.started":"2026-01-13T12:54:35.804177Z","shell.execute_reply":"2026-01-13T12:54:35.982845Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Example clean embedding ID:\", list(train_emb_clean.keys())[:3])\nprint(\"Example train_terms ID:\", train_terms[\"protein_id\"].unique()[:3])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:35.984688Z","iopub.execute_input":"2026-01-13T12:54:35.984864Z","iopub.status.idle":"2026-01-13T12:54:36.043412Z","shell.execute_reply.started":"2026-01-13T12:54:35.984848Z","shell.execute_reply":"2026-01-13T12:54:36.041999Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"reduced_data = {}\n\nfor ont in [\"F\", \"P\", \"C\"]:\n    subset = train_terms[\n        (train_terms[\"ontology\"] == ont) &\n        (train_terms[\"go_id\"].isin(top_go_terms[ont]))\n    ]\n\n    prot2labels = (\n        subset.groupby(\"protein_id\")[\"go_id\"]\n        .apply(set)\n        .to_dict()\n    )\n\n    X, Y = [], []\n\n    for pid, emb in train_emb_clean.items():\n        if pid in prot2labels:\n            X.append(emb)\n            Y.append(prot2labels[pid])\n\n    print(f\"{ont}: matched proteins = {len(X)}\")\n\n    if len(X) == 0:\n        print(f\"⚠️ No data for ontology {ont}, skipping\")\n        continue\n\n    mlb = MultiLabelBinarizer(classes=top_go_terms[ont])\n    Y_bin = mlb.fit_transform(Y)\n\n    reduced_data[ont] = {\n        \"X\": np.array(X),\n        \"Y\": Y_bin,\n        \"mlb\": mlb\n    }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:36.044318Z","iopub.execute_input":"2026-01-13T12:54:36.044544Z","iopub.status.idle":"2026-01-13T12:54:38.994123Z","shell.execute_reply.started":"2026-01-13T12:54:36.044526Z","shell.execute_reply":"2026-01-13T12:54:38.993033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models = {}\n\nfor ont in [\"F\", \"P\", \"C\"]:\n    print(f\"Training {ont} model...\")\n    \n    clf = OneVsRestClassifier(\n        SGDClassifier(\n            loss=\"log_loss\",\n            max_iter=5,        # SAFE\n            tol=1e-3,\n            n_jobs=-1\n        )\n    )\n    \n    clf.fit(\n        reduced_data[ont][\"X\"],\n        reduced_data[ont][\"Y\"]\n    )\n    \n    models[ont] = clf\n    print(f\"{ont} model done\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:54:38.995096Z","iopub.execute_input":"2026-01-13T12:54:38.995364Z","iopub.status.idle":"2026-01-13T13:02:03.141379Z","shell.execute_reply.started":"2026-01-13T12:54:38.995345Z","shell.execute_reply":"2026-01-13T13:02:03.140178Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"OUTPUT_PATH = \"/kaggle/working/submission.tsv\"\n\nTOP_K = 100\nCHUNK = 2000\n0\ntest_ids = list(test_emb.keys())\n\nwith open(OUTPUT_PATH, \"w\") as f:\n    for start in tqdm(range(0, len(test_ids), CHUNK)):\n        end = start + CHUNK\n        chunk_ids = test_ids[start:end]\n        X_chunk = np.array([test_emb[pid] for pid in chunk_ids])\n\n        for ont in models:\n            probs = models[ont].predict_proba(X_chunk)\n            mlb = reduced_data[ont][\"mlb\"]\n\n            for i, pid in enumerate(chunk_ids):\n                top_idx = probs[i].argsort()[-TOP_K:]\n                for idx in top_idx:\n                    score = probs[i][idx]\n                    if score > 0:\n                        f.write(\n                            f\"{pid}\\t{mlb.classes_[idx]}\\t{round(float(score),3)}\\n\"\n                        )\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T13:23:47.211855Z","iopub.execute_input":"2026-01-13T13:23:47.212741Z","iopub.status.idle":"2026-01-13T13:27:50.631677Z","shell.execute_reply.started":"2026-01-13T13:23:47.212712Z","shell.execute_reply":"2026-01-13T13:27:50.629635Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}