{"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":"gpu","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":"markdown","source":"This approach explores if Bert-based embeddings are helpful.","metadata":{}},{"cell_type":"code","source":"%pip install biopython","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:01.519006Z","iopub.execute_input":"2025-12-22T11:03:01.519596Z","iopub.status.idle":"2025-12-22T11:03:04.699625Z","shell.execute_reply.started":"2025-12-22T11:03:01.519560Z","shell.execute_reply":"2025-12-22T11:03:04.698761Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import transformers\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\nfrom tqdm import tqdm\nfrom Bio import SeqIO\nimport unicodedata\nimport re\nfrom joblib import Parallel, delayed\nfrom time import time\nfrom collections import defaultdict\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.neighbors import NearestNeighbors","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:04.701420Z","iopub.execute_input":"2025-12-22T11:03:04.701759Z","iopub.status.idle":"2025-12-22T11:03:09.870474Z","shell.execute_reply.started":"2025-12-22T11:03:04.701712Z","shell.execute_reply":"2025-12-22T11:03:09.869894Z"}},"outputs":[],"execution_count":null},{"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":"2025-12-22T11:03:09.871363Z","iopub.execute_input":"2025-12-22T11:03:09.871819Z","iopub.status.idle":"2025-12-22T11:03:09.883271Z","shell.execute_reply.started":"2025-12-22T11:03:09.871776Z","shell.execute_reply":"2025-12-22T11:03:09.882609Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Embedding Model","metadata":{}},{"cell_type":"code","source":"\n# Choose a CANINE variant; \"google/canine-c\" is the canonical character-level model\nMODEL_NAME = \"google/canine-c\"\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Device: {device}\")\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModel.from_pretrained(MODEL_NAME).to(device)\nmodel.eval()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:09.884995Z","iopub.execute_input":"2025-12-22T11:03:09.885287Z","iopub.status.idle":"2025-12-22T11:03:15.590705Z","shell.execute_reply.started":"2025-12-22T11:03:09.885256Z","shell.execute_reply":"2025-12-22T11:03:15.590070Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def canine_cls_embedding(texts):\n    \"\"\"\n    Returns a single embedding per input (the model's pooled [CLS] embedding).\n    \"\"\"\n\n    # Tokenize raw characters directly\n    enc = tokenizer(\n        texts,\n        padding=True,\n        truncation=True,\n        return_tensors=\"pt\"\n    ).to(device)\n\n    with torch.no_grad():\n        outputs = model(**enc)\n        # pooled_output is analogous to BERT's [CLS] embedding\n        # Shape: (batch_size, hidden_size)\n        cls_embeddings = outputs.pooler_output\n\n    # Return CPU numpy arrays\n    return cls_embeddings.detach().cpu().numpy()\n\n\ndef canine_mean_pool(texts):\n    \"\"\"\n    Returns one embedding per input text by mean-pooling the last hidden states\n    using the attention mask (ignoring padding).\n    \"\"\"\n\n    # Tokenize\n    enc = tokenizer(\n        texts,\n        padding=True,\n        truncation=True,\n        return_tensors=\"pt\"\n    ).to(device)\n\n    with torch.no_grad():\n        outputs = model(**enc)\n        # last_hidden_state: (batch_size, seq_len, hidden_size)\n        last_hidden = outputs.last_hidden_state  # torch.Tensor\n\n        # attention_mask: (batch_size, seq_len) -> expand to match hidden dims\n        attention_mask = enc[\"attention_mask\"].unsqueeze(-1)  # (batch, seq_len, 1)\n\n        # Mask out paddings\n        masked_hidden = last_hidden * attention_mask  # (batch, seq_len, hidden)\n\n        # Sum over sequence length and divide by the number of valid (non-pad) tokens\n        sum_hidden = masked_hidden.sum(dim=1)       # (batch, hidden)\n        lengths = attention_mask.sum(dim=1)         # (batch, 1)\n        mean_pooled = sum_hidden / torch.clamp(lengths, min=1e-9)\n\n    return mean_pooled.detach().cpu().numpy()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:15.591567Z","iopub.execute_input":"2025-12-22T11:03:15.592122Z","iopub.status.idle":"2025-12-22T11:03:15.598611Z","shell.execute_reply.started":"2025-12-22T11:03:15.592095Z","shell.execute_reply":"2025-12-22T11:03:15.597809Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# testing 1st method\ncanine_cls_embedding(['ABCD', 'EFGH'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:15.599514Z","iopub.execute_input":"2025-12-22T11:03:15.599883Z","iopub.status.idle":"2025-12-22T11:03:15.948324Z","shell.execute_reply.started":"2025-12-22T11:03:15.599860Z","shell.execute_reply":"2025-12-22T11:03:15.947675Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# testing 2nd method\ncanine_mean_pool(['ABCD', 'EFGH'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:15.949241Z","iopub.execute_input":"2025-12-22T11:03:15.949519Z","iopub.status.idle":"2025-12-22T11:03:15.997151Z","shell.execute_reply.started":"2025-12-22T11:03:15.949496Z","shell.execute_reply":"2025-12-22T11:03:15.996652Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# tokenization\nsequence = \"MRWQEMGYIFYPRKLR\"\n\ntokens = tokenizer.tokenize(sequence)\ntoken_ids = tokenizer.encode(sequence)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:15.998036Z","iopub.execute_input":"2025-12-22T11:03:15.998303Z","iopub.status.idle":"2025-12-22T11:03:16.002229Z","shell.execute_reply.started":"2025-12-22T11:03:15.998279Z","shell.execute_reply":"2025-12-22T11:03:16.001423Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tokens","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:16.003170Z","iopub.execute_input":"2025-12-22T11:03:16.003557Z","iopub.status.idle":"2025-12-22T11:03:16.016701Z","shell.execute_reply.started":"2025-12-22T11:03:16.003525Z","shell.execute_reply":"2025-12-22T11:03:16.016118Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Data Preparation","metadata":{}},{"cell_type":"code","source":"# load the train sequences\ntrain_sequence_file = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\ntotal = sum(1 for line in open(train_sequence_file, \"r\", encoding=\"utf-8\") if line.startswith(\">\"))\nprint(f\"Total number of train records: {total}\")\n\ntrain_sequence_df = []\nfor record in tqdm(SeqIO.parse(train_sequence_file, \"fasta\"), total=total, desc=\"Parsing FASTA\", unit=\"seq\"):\n    r = [\n        record.id, \n        record.id.split('|')[1],\n        record.description, \n        str(record.seq)\n    ]\n    train_sequence_df.append(r)\n    \n# convert to pandas data frame\ntrain_sequence_df = pd.DataFrame(train_sequence_df, columns=['full_prot_id', 'prot_id', 'description', 'sequence'])\ntrain_sequence_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:16.019117Z","iopub.execute_input":"2025-12-22T11:03:16.019524Z","iopub.status.idle":"2025-12-22T11:03:16.673361Z","shell.execute_reply.started":"2025-12-22T11:03:16.019499Z","shell.execute_reply":"2025-12-22T11:03:16.672797Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load train-terms\ntrain_term_df = pd.read_csv(\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n    sep=\"\\t\"\n)\ntrain_term_df = train_term_df.rename(columns={'EntryID': 'prot_id'})\ntrain_term_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:16.674271Z","iopub.execute_input":"2025-12-22T11:03:16.674533Z","iopub.status.idle":"2025-12-22T11:03:16.905141Z","shell.execute_reply.started":"2025-12-22T11:03:16.674509Z","shell.execute_reply":"2025-12-22T11:03:16.904240Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# merge to get labels\ntrain_sequence_df = train_sequence_df.merge(\n    train_term_df,\n    on='prot_id'\n)\ntrain_sequence_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:16.906415Z","iopub.execute_input":"2025-12-22T11:03:16.906702Z","iopub.status.idle":"2025-12-22T11:03:17.098293Z","shell.execute_reply.started":"2025-12-22T11:03:16.906676Z","shell.execute_reply":"2025-12-22T11:03:17.097540Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get the train series\n# note: two proteins can have the same sequence\ntrain = train_sequence_df.drop_duplicates(subset=['prot_id'])\ntrain.index = train.prot_id\ntrain = train['sequence']\nprint(f\"No. proteins: {train.index.nunique()}\")\nprint(f\"No. sequences: {train.drop_duplicates().shape[0]}\")\ntrain","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:17.099145Z","iopub.execute_input":"2025-12-22T11:03:17.099472Z","iopub.status.idle":"2025-12-22T11:03:17.491437Z","shell.execute_reply.started":"2025-12-22T11:03:17.099426Z","shell.execute_reply":"2025-12-22T11:03:17.490881Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load test data\ntest_superset_file = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\n# Count sequences quickly (without parsing)\ntotal = sum(1 for line in open(test_superset_file, \"r\", encoding=\"utf-8\") if line.startswith(\">\"))\nprint(f\"Total number of records: {total}\")\n\n# Load all sequences from a FASTA file\ntest_df = []\nfor record in tqdm(SeqIO.parse(test_superset_file, \"fasta\"), total=total, desc=\"Parsing FASTA\", unit=\"seq\"):\n    r = [record.id, record.description, str(record.seq)]\n    test_df.append(r)\n\ntest_df = pd.DataFrame(test_df, columns=['prot_id', 'description', 'sequence'])\nprint(f\"No. proteins: {test_df.prot_id.nunique()}\")\nprint(f\"No. sequences: {test_df.sequence.nunique()}\")\ntest_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:17.492466Z","iopub.execute_input":"2025-12-22T11:03:17.492912Z","iopub.status.idle":"2025-12-22T11:03:19.627582Z","shell.execute_reply.started":"2025-12-22T11:03:17.492885Z","shell.execute_reply":"2025-12-22T11:03:19.627003Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get test series\ntest = test_df.copy()\ntest.index = test.prot_id\ntest = test.sequence\ntest","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:19.628352Z","iopub.execute_input":"2025-12-22T11:03:19.628642Z","iopub.status.idle":"2025-12-22T11:03:19.666686Z","shell.execute_reply.started":"2025-12-22T11:03:19.628610Z","shell.execute_reply":"2025-12-22T11:03:19.666035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # testing only\n# train = train.iloc[:10000]\n# test = test.iloc[:1000]\n# test.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:19.667692Z","iopub.execute_input":"2025-12-22T11:03:19.668016Z","iopub.status.idle":"2025-12-22T11:03:19.682751Z","shell.execute_reply.started":"2025-12-22T11:03:19.667978Z","shell.execute_reply":"2025-12-22T11:03:19.682176Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# generate embeddings for train -> X_train\nt_start = time()\nBATCH_SIZE = 24\nn_batch_train = len(train)//BATCH_SIZE + 1\nX_train = []\nfor i in tqdm(range(n_batch_train)):\n    start = i*BATCH_SIZE\n    end = min((i+1)*BATCH_SIZE, len(train))\n    batch_texts = train.iloc[start: end].tolist()\n    if len(batch_texts) == 0:\n        continue\n    batch_emb = canine_mean_pool(batch_texts)\n    X_train.append(batch_emb)\nX_train = np.concatenate(X_train)\nassert X_train.shape[0] == len(train), \"missing texts!\"\nprint(f\"Train embeddings extracted in {time() - t_start} seconds!\")\nX_train.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:03:19.683598Z","iopub.execute_input":"2025-12-22T11:03:19.683889Z","iopub.status.idle":"2025-12-22T11:08:26.916351Z","shell.execute_reply.started":"2025-12-22T11:03:19.683853Z","shell.execute_reply":"2025-12-22T11:08:26.915779Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# generate embeddings for test -> X_test\nt_start = time()\nn_batch_test = len(test)//BATCH_SIZE + 1\nX_test = []\nfor i in tqdm(range(n_batch_test)):\n    start = i*BATCH_SIZE\n    end = min((i+1)*BATCH_SIZE, len(test))\n    batch_texts = test.iloc[start: end].tolist()\n    if len(batch_texts) == 0:\n        continue\n    batch_emb = canine_mean_pool(batch_texts)\n    X_test.append(batch_emb)\nX_test = np.concatenate(X_test)\nassert len(X_test) == len(test), \"missing texts in test!\"\nprint(f\"Test embeddings extracted in {time() - t_start} seconds!\")\nX_test.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:08:26.917432Z","iopub.execute_input":"2025-12-22T11:08:26.918078Z","iopub.status.idle":"2025-12-22T11:08:57.983799Z","shell.execute_reply.started":"2025-12-22T11:08:26.918052Z","shell.execute_reply":"2025-12-22T11:08:57.983063Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Similarity Search","metadata":{}},{"cell_type":"code","source":"# KNN\nprint(f\"fitting KNN...\")\nt_start = time()\nnn = NearestNeighbors(\n    n_neighbors=3,\n    metric='cosine',       # change to 'euclidean' if you prefer\n    algorithm='brute',     # brute is reliable for sparse + cosine\n    n_jobs=-1              # use all cores if available\n)\n\n# Fit on the training set\nnn.fit(X_train)\nprint(f\"KNN fitted in {time() - t_start}\")\n\nprint(f\"querrying...\")\nt_start = time()\n\n# Chunked inference\nn_test = X_test.shape[0]\nk = 3\nbatch_size = 2000  # tune based on your RAM/CPU\nall_indices = np.empty((n_test, k), dtype=np.int32)\nall_distances = np.empty((n_test, k), dtype=np.float32)\n\nfor start in tqdm(range(0, n_test, batch_size)):\n    end = min(start + batch_size, n_test)\n    X_batch = X_test[start:end]\n    distances, indices = nn.kneighbors(X_batch, n_neighbors=k, return_distance=True)\n    all_indices[start:end] = indices\n    all_distances[start:end] = distances\n\n# Optionally convert to similarities\nall_similarities = 1.0 - all_distances\n\nprint(f\"finished searching in {time() - t_start} seconds\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:08:57.984853Z","iopub.execute_input":"2025-12-22T11:08:57.985151Z","iopub.status.idle":"2025-12-22T11:08:58.231638Z","shell.execute_reply.started":"2025-12-22T11:08:57.985119Z","shell.execute_reply":"2025-12-22T11:08:58.229667Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# collect labels\ntest_ids = np.repeat(np.arange(n_test, dtype=np.int32), k)\ntrain_ids = all_indices.ravel().astype(np.int32)\nsimilarity_scores = all_similarities.ravel().astype(np.float32)\n\nknn_df = pd.DataFrame({\n    \"test_id\": test_ids,\n    \"train_id\": train_ids,\n    \"similarity_score\": similarity_scores\n})\n\ntest_ind2id = dict(zip(range(test.shape[0]), test.index))\ntrain_ind2id = dict(zip(range(train.shape[0]), train.index))\n\nknn_df['test_prot'] = knn_df.test_id.map(test_ind2id)\nknn_df['train_prot'] = knn_df.train_id.map(train_ind2id)\nknn_df = knn_df.merge(\n    train_sequence_df[['prot_id', 'term']],\n    left_on='train_prot',\n    right_on='prot_id',\n    how='left'\n)\nknn_df = knn_df[['test_prot', 'term', 'similarity_score']]\nknn_df = knn_df\\\n    .sort_values(\n        by=['test_prot', 'term', 'similarity_score'], \n        ascending=[True, True, False])\\\n    .drop_duplicates(\n        subset=['test_prot', 'term'],\n        keep='first'\n    ).reset_index(drop=True)\nknn_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:08:58.232751Z","iopub.execute_input":"2025-12-22T11:08:58.232998Z","iopub.status.idle":"2025-12-22T11:08:58.364089Z","shell.execute_reply.started":"2025-12-22T11:08:58.232972Z","shell.execute_reply":"2025-12-22T11:08:58.363304Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# save to file\nknn_df.to_csv('submission.tsv', index=False, sep='\\t', header=False)\nknn_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T11:08:58.365065Z","iopub.execute_input":"2025-12-22T11:08:58.365718Z","iopub.status.idle":"2025-12-22T11:08:58.427053Z","shell.execute_reply.started":"2025-12-22T11:08:58.365693Z","shell.execute_reply":"2025-12-22T11:08:58.426479Z"}},"outputs":[],"execution_count":null}]}