{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31236,"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":"2026-01-11T21:32:19.638371Z","iopub.execute_input":"2026-01-11T21:32:19.638574Z","iopub.status.idle":"2026-01-11T21:32:20.681031Z","shell.execute_reply.started":"2026-01-11T21:32:19.638552Z","shell.execute_reply":"2026-01-11T21:32:20.680311Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install biopython transformers --quiet\n!pip install sentencepiece --quiet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:32:20.682742Z","iopub.execute_input":"2026-01-11T21:32:20.683255Z","iopub.status.idle":"2026-01-11T21:32:30.244907Z","shell.execute_reply.started":"2026-01-11T21:32:20.683227Z","shell.execute_reply":"2026-01-11T21:32:30.243979Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import libraries\nimport pandas as pd\nimport numpy as np\nfrom Bio import SeqIO\nimport torch\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:32:30.246381Z","iopub.execute_input":"2026-01-11T21:32:30.247124Z","iopub.status.idle":"2026-01-11T21:32:33.808641Z","shell.execute_reply.started":"2026-01-11T21:32:30.247091Z","shell.execute_reply":"2026-01-11T21:32:33.808035Z"}},"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\ntest_seqs = seqs(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:32:33.809527Z","iopub.execute_input":"2026-01-11T21:32:33.809964Z","iopub.status.idle":"2026-01-11T21:32:35.621177Z","shell.execute_reply.started":"2026-01-11T21:32:33.809928Z","shell.execute_reply":"2026-01-11T21:32:35.620485Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_seqs.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:32:35.622115Z","iopub.execute_input":"2026-01-11T21:32:35.622414Z","iopub.status.idle":"2026-01-11T21:32:35.646490Z","shell.execute_reply.started":"2026-01-11T21:32:35.622380Z","shell.execute_reply":"2026-01-11T21:32:35.645617Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_seqs.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:32:35.647609Z","iopub.execute_input":"2026-01-11T21:32:35.647933Z","iopub.status.idle":"2026-01-11T21:32:35.653311Z","shell.execute_reply.started":"2026-01-11T21:32:35.647906Z","shell.execute_reply":"2026-01-11T21:32:35.652638Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load ProtT5 model and tokenizer\nimport 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":"2026-01-11T21:32:57.040556Z","iopub.execute_input":"2026-01-11T21:32:57.041267Z","iopub.status.idle":"2026-01-11T21:33:21.095595Z","shell.execute_reply.started":"2026-01-11T21:32:57.041234Z","shell.execute_reply":"2026-01-11T21:33:21.094800Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tokenizer = 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":"2026-01-11T21:33:36.795591Z","iopub.execute_input":"2026-01-11T21:33:36.796355Z","iopub.status.idle":"2026-01-11T21:34:59.190507Z","shell.execute_reply.started":"2026-01-11T21:33:36.796322Z","shell.execute_reply":"2026-01-11T21:34:59.189825Z"}},"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:36:13.015906Z","iopub.execute_input":"2026-01-11T21:36:13.016401Z","iopub.status.idle":"2026-01-11T21:36:13.024007Z","shell.execute_reply.started":"2026-01-11T21:36:13.016373Z","shell.execute_reply":"2026-01-11T21:36:13.023167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# test embeddings\nall_embeddings = []\n\nfor seq in tqdm(test_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_test_embeddings.pt\")\nnp.save(\"/kaggle/working/protT5_test_embeddings.npy\", all_embeddings)\n\nprint(\"Embedding completed.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T21:39:04.272414Z","iopub.execute_input":"2026-01-11T21:39:04.272792Z"}},"outputs":[],"execution_count":null}]}