{"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":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":13991224,"sourceType":"datasetVersion","datasetId":8917191}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Logistic Regression with ESM-2 embeddings\n\nIdea:\n- Model: OneVsRestClassifier(LogisticRegression) --> train one Logistic Regression model for one class\n- Features: PCA(n_components=100) --> PCA.fit_transform(ESM-2 embeddings)\n- Labels: Three sets for three ontologies (P, C, F)\n    - P has 16858 classes\n    - C has 2651 classes\n    - F has 6616 classes\n\nFuture direction:\n- Model --> MLP\n- Features --> full embeddings\n\nReferences:\n- (EDA + OneVsRestClassifier) https://www.kaggle.com/code/analyticaobscura/cafa-6-decoding-protein-mysteries\n- (ESM-2 320-D embeddings) https://www.kaggle.com/code/dalloliogm/compute-protein-embeddings-with-esm2-esm-c/notebook\n- (Optional ProtT5 1024-D embeddings) https://www.kaggle.com/code/ahsuna123/t5-embedding-calculation-cafa-6/output?select=train_ids.npy\n\n---","metadata":{}},{"cell_type":"code","source":"!pip install biopython > /dev/null","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:00.755971Z","iopub.execute_input":"2025-12-07T03:25:00.756274Z","iopub.status.idle":"2025-12-07T03:25:07.400674Z","shell.execute_reply.started":"2025-12-07T03:25:00.756240Z","shell.execute_reply":"2025-12-07T03:25:07.399534Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 1: Load CAFA6 files\n\n---","metadata":{}},{"cell_type":"code","source":"# CAFA6 file paths\nTRAIN_TERMS = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\nTRAIN_SEQ = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\nTEST_SEQ = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:07.403026Z","iopub.execute_input":"2025-12-07T03:25:07.403330Z","iopub.status.idle":"2025-12-07T03:25:07.408047Z","shell.execute_reply.started":"2025-12-07T03:25:07.403302Z","shell.execute_reply":"2025-12-07T03:25:07.407144Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio import SeqIO \n\n# Dict {entryId, seq}\ntrain_sequences = {rec.id: str(rec.seq) for rec in SeqIO.parse(TRAIN_SEQ, 'fasta')}\ntest_sequences  = {rec.id: str(rec.seq) for rec in SeqIO.parse(TEST_SEQ,  'fasta')}\n\nprint(f'Loaded {len(train_sequences)} train and {len(test_sequences)} test sequences')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:08.702265Z","iopub.execute_input":"2025-12-07T03:25:08.703380Z","iopub.status.idle":"2025-12-07T03:25:12.046622Z","shell.execute_reply.started":"2025-12-07T03:25:08.703338Z","shell.execute_reply":"2025-12-07T03:25:12.045859Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Train dict:\", list(train_sequences.items())[0])\nprint(\"Test dict:\", list(test_sequences.items())[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:12.047976Z","iopub.execute_input":"2025-12-07T03:25:12.048242Z","iopub.status.idle":"2025-12-07T03:25:12.125924Z","shell.execute_reply.started":"2025-12-07T03:25:12.048218Z","shell.execute_reply":"2025-12-07T03:25:12.125144Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ids = [i.split('|')[1] for i in train_sequences.keys()]\ntest_ids = list(test_sequences.keys())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:12.126885Z","iopub.execute_input":"2025-12-07T03:25:12.127286Z","iopub.status.idle":"2025-12-07T03:25:12.153206Z","shell.execute_reply.started":"2025-12-07T03:25:12.127262Z","shell.execute_reply":"2025-12-07T03:25:12.152243Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"train_ids[0:10]:\", train_ids[0:10])\nprint(\"test_ids[0:10]:\", test_ids[0:10])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:12.154762Z","iopub.execute_input":"2025-12-07T03:25:12.155238Z","iopub.status.idle":"2025-12-07T03:25:12.172756Z","shell.execute_reply.started":"2025-12-07T03:25:12.155215Z","shell.execute_reply":"2025-12-07T03:25:12.172038Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 2: Feature extraction\n\n---","metadata":{}},{"cell_type":"code","source":"# Embeddings file paths\nESM_EMBEDDINGS = \"/kaggle/input/esm2-t6-8m-ur50d-cafa6\"\nTRAIN_EMBEDDINGS = ESM_EMBEDDINGS + \"/protein_embeddings_train.npy\"\nTEST_EMBEDDINGS = ESM_EMBEDDINGS + \"/protein_embeddings_test.npy\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:14.417593Z","iopub.execute_input":"2025-12-07T03:25:14.418421Z","iopub.status.idle":"2025-12-07T03:25:14.422545Z","shell.execute_reply.started":"2025-12-07T03:25:14.418395Z","shell.execute_reply":"2025-12-07T03:25:14.421424Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\n# Load embeddings\nX_train = np.load(TRAIN_EMBEDDINGS)\nX_test = np.load(TEST_EMBEDDINGS)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:14.692534Z","iopub.execute_input":"2025-12-07T03:25:14.692837Z","iopub.status.idle":"2025-12-07T03:25:17.290869Z","shell.execute_reply.started":"2025-12-07T03:25:14.692816Z","shell.execute_reply":"2025-12-07T03:25:17.289925Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"X_train shape:\", X_train.shape)\nprint(\"X_test shape:\", X_test.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:17.292326Z","iopub.execute_input":"2025-12-07T03:25:17.292934Z","iopub.status.idle":"2025-12-07T03:25:17.297764Z","shell.execute_reply.started":"2025-12-07T03:25:17.292907Z","shell.execute_reply":"2025-12-07T03:25:17.296991Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.decomposition import PCA\n\npca = PCA(n_components=100, random_state=42)\nX_train_reduced = pca.fit_transform(X_train)\nX_test_reduced  = pca.transform(X_test)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:17.298818Z","iopub.execute_input":"2025-12-07T03:25:17.299659Z","iopub.status.idle":"2025-12-07T03:25:20.846289Z","shell.execute_reply.started":"2025-12-07T03:25:17.299636Z","shell.execute_reply":"2025-12-07T03:25:20.845303Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"X_train_reduced shape:\", X_train_reduced.shape)\nprint(\"X_test_reduced shape:\", X_test_reduced.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:20.847860Z","iopub.execute_input":"2025-12-07T03:25:20.848327Z","iopub.status.idle":"2025-12-07T03:25:20.853580Z","shell.execute_reply.started":"2025-12-07T03:25:20.848294Z","shell.execute_reply":"2025-12-07T03:25:20.852521Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 3: Label encoding and Training\n\n---","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom tqdm import tqdm\n\nmlb_dict = dict()\nmodels = dict()\n\ntrain_terms_df = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\n\nfor aspect in tqdm(['P', 'C', 'F'], desc=\"Training models\"):\n    # Filter the train_terms_df based on aspect\n    ont_terms_df = train_terms_df[train_terms_df['aspect'] == aspect]\n\n    # Group the dataFrame based on the EntryID, turn all the GO terms to a list, finally turns this dataFrame to a dict {entryID: terms}\n    protein_terms = ont_terms_df.groupby('EntryID')['term'].apply(list).to_dict()\n\n    # Create a list of labels for this aspect, if an entryID doesn't exist in this aspect, give it a []\n    # This ensures y_train is of shape (82404, ...)\n    labels = [protein_terms.get(entry_id, []) for entry_id in train_ids]\n\n    # Multi-hot encoding, use sparse representation\n    mlb = MultiLabelBinarizer(sparse_output=True)\n    y_train = mlb.fit_transform(labels)\n    \n    print(f\"y_train shape for {aspect} ontology: {y_train.shape} \\t Number of unique {aspect} terms: {y_train.shape[1]}\")\n\n    # Save to dict\n    mlb_dict[aspect] = mlb\n    model = OneVsRestClassifier(LogisticRegression(max_iter=600, solver='lbfgs', C=0.5, random_state=42), n_jobs=-1)\n    model.fit(X_train_reduced, y_train)\n    models[aspect] = model\n    print(f\"Model for {aspect} trained successfully.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:25:23.180037Z","iopub.execute_input":"2025-12-07T03:25:23.181034Z","iopub.status.idle":"2025-12-07T03:25:27.935856Z","shell.execute_reply.started":"2025-12-07T03:25:23.181005Z","shell.execute_reply":"2025-12-07T03:25:27.935082Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 4: Inference and Submission","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 5000  # avoid memory overflow\nsubmission_list = []\n\nfor i in tqdm(range(0, len(test_ids), BATCH_SIZE), desc=\"Predicting on Test Set\"):\n    batch_entry_ids = test_ids[i : i + BATCH_SIZE]\n\n    X_batch = X_test_reduced[i : i + BATCH_SIZE]\n\n    for aspect, model in models.items():\n        mlb = mlb_dict[aspect]\n        y_pred_proba = model.predict_proba(X_batch)\n\n        for j, entry_id in enumerate(batch_entry_ids):\n            probs = y_pred_proba[j]\n            candidate_indices = np.where(probs > 0.02)[0]\n\n            for idx in candidate_indices:\n                submission_list.append((entry_id, mlb.classes_[idx], round(probs[idx], 3)))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df = pd.DataFrame(submission_list, columns=['Protein Id', 'GO Term Id', 'Prediction'])\n\nprint(\"Applying 1500 prediction limit per protein...\")\nsubmission_df = submission_df.sort_values(by=['Protein Id', 'Prediction'], ascending=[True, False])\nfinal_submission_df = submission_df.groupby('Protein Id').head(1500).reset_index(drop=True)\nfinal_submission_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)\n\nprint(\"\\nSubmission file 'submission.tsv' created successfully.\")\nprint(f\"Total predictions in final submission: {len(final_submission_df):,}\")\nprint(\"Submission DataFrame Head:\")\ndisplay(final_submission_df.head())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---\n\n- Result: 5M2 predictions\n- Score: 0.230\n- Upvote if helpful","metadata":{}}]}