{"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":14084779,"sourceType":"competition"}],"dockerImageVersionId":31193,"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":"2025-12-07T18:06:13.927591Z","iopub.execute_input":"2025-12-07T18:06:13.927892Z","iopub.status.idle":"2025-12-07T18:06:14.198146Z","shell.execute_reply.started":"2025-12-07T18:06:13.927868Z","shell.execute_reply":"2025-12-07T18:06:14.197433Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Laod the sequence data","metadata":{}},{"cell_type":"code","source":"!pip install biopython transformers --quiet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:22.7754Z","iopub.execute_input":"2025-12-07T18:06:22.775801Z","iopub.status.idle":"2025-12-07T18:06:28.348294Z","shell.execute_reply.started":"2025-12-07T18:06:22.775777Z","shell.execute_reply":"2025-12-07T18:06:28.347362Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install sentencepiece","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:30.827457Z","iopub.execute_input":"2025-12-07T18:06:30.827762Z","iopub.status.idle":"2025-12-07T18:06:34.05527Z","shell.execute_reply.started":"2025-12-07T18:06:30.827735Z","shell.execute_reply":"2025-12-07T18:06:34.054295Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import libraries\nimport pandas as pd\nimport numpy as np\nfrom Bio import SeqIO\nimport torch\nprint(\"installation completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:35.699705Z","iopub.execute_input":"2025-12-07T18:06:35.700596Z","iopub.status.idle":"2025-12-07T18:06:39.176007Z","shell.execute_reply.started":"2025-12-07T18:06:35.700559Z","shell.execute_reply":"2025-12-07T18:06:39.175289Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seqs(fasta_path):\n    records = []\n    for record in SeqIO.parse(fasta_path,\"fasta\"):\n        records.append({\n            \"id\":record.id,\n            \"sequence\": str(record.seq)\n        })\n\n    return pd.DataFrame(records)\n\ntrain_seqs = seqs(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\")\ntest_seqs = seqs(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:41.579961Z","iopub.execute_input":"2025-12-07T18:06:41.580909Z","iopub.status.idle":"2025-12-07T18:06:44.325613Z","shell.execute_reply.started":"2025-12-07T18:06:41.580879Z","shell.execute_reply":"2025-12-07T18:06:44.324995Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seqs.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:43:55.050899Z","iopub.execute_input":"2025-12-07T03:43:55.051181Z","iopub.status.idle":"2025-12-07T03:43:55.07585Z","shell.execute_reply.started":"2025-12-07T03:43:55.051161Z","shell.execute_reply":"2025-12-07T03:43:55.07498Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seqs.shape[0]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:43:58.402618Z","iopub.execute_input":"2025-12-07T03:43:58.403258Z","iopub.status.idle":"2025-12-07T03:43:58.407927Z","shell.execute_reply.started":"2025-12-07T03:43:58.403231Z","shell.execute_reply":"2025-12-07T03:43:58.407279Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_seqs.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T03:44:03.548702Z","iopub.execute_input":"2025-12-07T03:44:03.549009Z","iopub.status.idle":"2025-12-07T03:44:03.557497Z","shell.execute_reply.started":"2025-12-07T03:44:03.548987Z","shell.execute_reply":"2025-12-07T03:44:03.556867Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Clean IDs\ndef seq_id_clean(raw_seqs):\n    raw_seqs[\"id\"] = raw_seqs[\"id\"].str.split(\"|\").str[1].fillna(raw_seqs[\"id\"])\n    \n    return raw_seqs\n    \ntrain_seqs = seq_id_clean(train_seqs)\ntrain_seqs.head()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:50.155291Z","iopub.execute_input":"2025-12-07T18:06:50.155586Z","iopub.status.idle":"2025-12-07T18:06:50.396831Z","shell.execute_reply.started":"2025-12-07T18:06:50.155563Z","shell.execute_reply":"2025-12-07T18:06:50.396174Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load GO tern and aspects","metadata":{}},{"cell_type":"code","source":"train_terms = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\", sep=\"\\t\")\ntrain_terms.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:53.446558Z","iopub.execute_input":"2025-12-07T18:06:53.446865Z","iopub.status.idle":"2025-12-07T18:06:53.751581Z","shell.execute_reply.started":"2025-12-07T18:06:53.446825Z","shell.execute_reply":"2025-12-07T18:06:53.750869Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms = train_terms.rename(columns = {\"EntryID\":\"id\"})\ntrain_terms.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:56.003532Z","iopub.execute_input":"2025-12-07T18:06:56.003827Z","iopub.status.idle":"2025-12-07T18:06:56.028584Z","shell.execute_reply.started":"2025-12-07T18:06:56.003805Z","shell.execute_reply":"2025-12-07T18:06:56.027765Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"go_to_aspect = {\n    row[\"term\"]: row[\"aspect\"]\n    for _, row in train_terms.dropna(subset=[\"term\", \"aspect\"]).iterrows()\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:06:58.823212Z","iopub.execute_input":"2025-12-07T18:06:58.823501Z","iopub.status.idle":"2025-12-07T18:07:18.713125Z","shell.execute_reply.started":"2025-12-07T18:06:58.823481Z","shell.execute_reply":"2025-12-07T18:07:18.712469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from scipy.sparse import csr_matrix\nimport numpy as np","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:07:33.103029Z","iopub.execute_input":"2025-12-07T18:07:33.103325Z","iopub.status.idle":"2025-12-07T18:07:33.192415Z","shell.execute_reply.started":"2025-12-07T18:07:33.103303Z","shell.execute_reply":"2025-12-07T18:07:33.191774Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# protein index mapping (in the same order as train_seqs)\nprotein_ids = train_seqs[\"id\"].tolist()\npid_to_idx = {pid: i for i, pid in enumerate(protein_ids)}\n\n# GO index mapping\ngo_terms = sorted(train_terms[\"term\"].unique())\ngo_to_idx = {go: i for i, go in enumerate(go_terms)}\n\n# build coordinate lists\nrow_idx = []\ncol_idx = []\n\nfor pid, go in zip(train_terms[\"id\"], train_terms[\"term\"]):\n    row_idx.append(pid_to_idx[pid])\n    col_idx.append(go_to_idx[go])\n\ndata = np.ones(len(row_idx), dtype=np.float32)\n\nY_sparse = csr_matrix(\n    (data, (row_idx, col_idx)),\n    shape=(len(protein_ids), len(go_terms)),\n    dtype=np.float32\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:07:35.862343Z","iopub.execute_input":"2025-12-07T18:07:35.863106Z","iopub.status.idle":"2025-12-07T18:07:36.388496Z","shell.execute_reply.started":"2025-12-07T18:07:35.863081Z","shell.execute_reply":"2025-12-07T18:07:36.387607Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\n\nY_coo = Y_sparse.tocoo()\n\nindices = torch.tensor(\n    np.vstack((Y_coo.row, Y_coo.col)),\n    dtype=torch.int64\n)\n\nvalues = torch.tensor(Y_coo.data, dtype=torch.float32)\n\nY_torch_sparse = torch.sparse_coo_tensor(\n    indices,\n    values,\n    size=Y_coo.shape,\n    dtype=torch.float32\n).coalesce()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:07:43.458642Z","iopub.execute_input":"2025-12-07T18:07:43.459162Z","iopub.status.idle":"2025-12-07T18:07:43.512939Z","shell.execute_reply.started":"2025-12-07T18:07:43.459139Z","shell.execute_reply":"2025-12-07T18:07:43.51213Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"Y_torch_sparse.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:07:46.123359Z","iopub.execute_input":"2025-12-07T18:07:46.123637Z","iopub.status.idle":"2025-12-07T18:07:46.12921Z","shell.execute_reply.started":"2025-12-07T18:07:46.123617Z","shell.execute_reply":"2025-12-07T18:07:46.128249Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load ProtT5 model and tokenizer","metadata":{}},{"cell_type":"markdown","source":"# ProtT5 requires:\n* Uppercase sequences\n* Replace uncommon amino acids with X\n* Add spaces between amino acids ","metadata":{}},{"cell_type":"code","source":"import torch\nfrom transformers import T5Tokenizer,T5EncoderModel\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:07:49.346525Z","iopub.execute_input":"2025-12-07T18:07:49.346809Z","iopub.status.idle":"2025-12-07T18:08:15.717264Z","shell.execute_reply.started":"2025-12-07T18:07:49.346786Z","shell.execute_reply":"2025-12-07T18:08:15.716553Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ntokenizer = T5Tokenizer.from_pretrained(\n    \"Rostlab/prot_t5_xl_uniref50\",\n    do_lower_case=False\n)\nmodel = T5EncoderModel.from_pretrained(\n    \"Rostlab/prot_t5_xl_uniref50\"\n).to(device).half()\nmodel.eval()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:08:19.955825Z","iopub.execute_input":"2025-12-07T18:08:19.956465Z","iopub.status.idle":"2025-12-07T18:09:28.399005Z","shell.execute_reply.started":"2025-12-07T18:08:19.956438Z","shell.execute_reply":"2025-12-07T18:09:28.398045Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Preprocessing sequences\ndef preprocess_seq(seq):\n    seq = seq.upper()\n    seq = seq.replace(\"U\",\"X\").replace(\"Z\",\"X\").replace(\"0\",\"X\").replace(\"B\",\"X\")\n    return \" \".join(list(seq))\n\ndef chunk_sequence(seq,max_len =510):\n    \"\"\"split sequence into 510-AA chunks if too long.\"\"\"\n    return [seq[i:i+max_len] for i in range(0,len(seq),max_len)]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:09:54.35719Z","iopub.execute_input":"2025-12-07T18:09:54.357462Z","iopub.status.idle":"2025-12-07T18:10:00.33573Z","shell.execute_reply.started":"2025-12-07T18:09:54.357439Z","shell.execute_reply":"2025-12-07T18:10:00.334741Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nfrom tqdm import tqdm\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n\n# preprocessing and chunking functions\n\ndef preprocess_seq(seq):\n    seq = seq.strip()\n    return seq\n\ndef chunk_sequence(seq, chunk_size=1022):\n    return [seq[i:i+chunk_size] for i in range(0, len(seq), chunk_size)]\n\n\ndef embed_one(seq):\n    seq = preprocess_seq(seq)\n    raw_seq = seq.replace(\"U\", \"X\").replace(\"Z\", \"X\").replace(\"B\", \"X\")\n\n    chunks = chunk_sequence(raw_seq)\n\n    chunk_embs = []\n    for c in chunks:\n        c_spaced = \" \".join(list(c))  \n        tokens = tokenizer(c_spaced, return_tensors=\"pt\").to(device)\n\n        with torch.no_grad():\n            out = model(**tokens).last_hidden_state\n\n        mask = tokens.attention_mask\n        emb = (out * mask.unsqueeze(-1)).sum(dim=1) / mask.sum(dim=1).unsqueeze(-1)\n        chunk_embs.append(emb.squeeze(0))\n\n    # average all chunk embeddings\n    final_emb = torch.stack(chunk_embs, dim=0).mean(dim=0)\n    return final_emb.cpu().numpy()\n\n\nall_embeddings = []\n\nfor seq in tqdm(train_seqs[\"sequence\"], desc=\"Embedding sequences\"):\n    emb = embed_one(seq)  \n    all_embeddings.append(emb)\n\nall_embeddings = np.vstack(all_embeddings)\n\n\n# Save outputs\n\ntorch.save(all_embeddings, \"/kaggle/working/protT5_embeddings.pt\")\nnp.save(\"/kaggle/working/protT5_embeddings.npy\", all_embeddings)\n\nprint(\"Embedding completed.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-07T18:10:05.956632Z","iopub.execute_input":"2025-12-07T18:10:05.957198Z","iopub.status.idle":"2025-12-07T22:43:21.500343Z","shell.execute_reply.started":"2025-12-07T18:10:05.957174Z","shell.execute_reply":"2025-12-07T22:43:21.499443Z"}},"outputs":[],"execution_count":null}]}