{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","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":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Protein Embedding Generator for CAFA6\n\nThis notebook generates embeddings for protein sequences using ESM2 or ESM-Cambrian models.\n\nESM-Cambrian is the replacement for ESM2 - more info here https://www.evolutionaryscale.ai/blog/esm-cambrian\n\nThere is no submission code here - just computing the embeddings.\n\nOutputs: `protein_embeddings_train.npy`, `protein_embeddings_test.npy`, and ID lists.","metadata":{}},{"cell_type":"code","source":"!uv pip install biopython\n!uv pip install -q \"protobuf<5\"\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:41.725113Z","iopub.execute_input":"2025-11-30T21:23:41.725740Z","iopub.status.idle":"2025-11-30T21:23:42.234402Z","shell.execute_reply.started":"2025-11-30T21:23:41.725709Z","shell.execute_reply":"2025-11-30T21:23:42.233549Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import google.protobuf\nprint(\"protobuf version:\", google.protobuf.__version__)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"MODEL_TYPE ='esm3'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:42.235857Z","iopub.execute_input":"2025-11-30T21:23:42.236107Z","iopub.status.idle":"2025-11-30T21:23:42.240243Z","shell.execute_reply.started":"2025-11-30T21:23:42.236083Z","shell.execute_reply":"2025-11-30T21:23:42.239453Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nfrom Bio import SeqIO\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ------------------\n# 1. Config\n# ------------------\nMODEL_TYPE = 'esm2'  # 'esm2' or 'esm3'\nESM2_MODEL_NAME = 'facebook/esm2_t6_8M_UR50D'\n# For ESM3 you'll need to confirm the exact model string & API\nESM3_MODEL_NAME = 'esm3_sm_open_v1'  \n\nBATCH_SIZE = 32        # try increasing while GPU memory allows\nMAX_LENGTH = 1022      # typical ESM2 max (excluding special tokens)\n\ndata_dir = '/kaggle/input/cafa-6-protein-function-prediction/Train/'\ntrain_seq_file = os.path.join(data_dir, 'train_sequences.fasta')\ntest_seq_file = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f'Using device: {device}')\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:42.240963Z","iopub.execute_input":"2025-11-30T21:23:42.241202Z","iopub.status.idle":"2025-11-30T21:23:42.256281Z","shell.execute_reply.started":"2025-11-30T21:23:42.241178Z","shell.execute_reply":"2025-11-30T21:23:42.255608Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------\n# 2. Load sequences\n# ------------------\n# If Debug mode is on, get just a subset of data\ndebug = False\nmax_train = 300\nmax_test = 10\n\nif debug:\n    train_sequences = {}\n    for i, rec in enumerate(SeqIO.parse(train_seq_file, 'fasta')):\n        if i >= max_train:\n            break\n        train_sequences[rec.id] = str(rec.seq)\n\n    test_sequences = {}\n    for i, rec in enumerate(SeqIO.parse(test_seq_file, 'fasta')):\n        if i >= max_test:\n            break\n        test_sequences[rec.id] = str(rec.seq)\n\nelse:\n    train_sequences = {rec.id: str(rec.seq) for rec in SeqIO.parse(train_seq_file, 'fasta')}\n    test_sequences  = {rec.id: str(rec.seq) for rec in SeqIO.parse(test_seq_file,  'fasta')}\n\nprint(f'Loaded {len(train_sequences)} train and {len(test_sequences)} test sequences')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:42.258121Z","iopub.execute_input":"2025-11-30T21:23:42.258417Z","iopub.status.idle":"2025-11-30T21:23:42.280111Z","shell.execute_reply.started":"2025-11-30T21:23:42.258400Z","shell.execute_reply":"2025-11-30T21:23:42.279426Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------\n# 3. Load model\n# ------------------\nif MODEL_TYPE == 'esm2':\n    tokenizer = AutoTokenizer.from_pretrained(ESM2_MODEL_NAME)\n    model = AutoModel.from_pretrained(ESM2_MODEL_NAME)\n    model.to(device)\n    model.eval()\n    print(f'Loaded ESM2 model: {ESM2_MODEL_NAME}')\n\nelif MODEL_TYPE == 'esm3':   # using ESM-Cambrian instead\n    import esm\n    # Choose model — adjust based on Kaggle GPU memory\n    model, alphabet = esm.pretrained.esm_camb_xl_t33_650M_UR50D()\n\n    batch_converter = alphabet.get_batch_converter()\n    model = model.to(device)\n    model.eval()\n\n    LAYER = 33   # final layer of the xl_t33 model\n\n    @torch.no_grad()\n    def get_cambrian_embeddings(batch_seqs):\n        # ESM Cambrian expects a list of tuples: (id, seq)\n        data = [(str(i), seq) for i, seq in enumerate(batch_seqs)]\n        batch_labels, batch_strs, batch_tokens = batch_converter(data)\n\n        batch_tokens = batch_tokens.to(device)\n\n        # Forward pass\n        out = model(batch_tokens, repr_layers=[LAYER])\n        reps = out[\"representations\"][LAYER]  # [batch, len, hidden]\n\n        # Remove start/end tokens\n        reps = reps[:, 1:-1, :]\n\n        # Mean pool over residues\n        emb = reps.mean(1).cpu().numpy()\n        return emb\n\n\n\nelse:\n    raise ValueError(\"MODEL_TYPE must be 'esm2' or 'esm3'\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:42.280778Z","iopub.execute_input":"2025-11-30T21:23:42.280955Z","iopub.status.idle":"2025-11-30T21:23:42.718035Z","shell.execute_reply.started":"2025-11-30T21:23:42.280936Z","shell.execute_reply":"2025-11-30T21:23:42.717319Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ------------------\n# 4. Helper: compute batch embeddings\n# ------------------\n@torch.no_grad()\ndef compute_embeddings(sequences, ids, batch_size=BATCH_SIZE):\n    \"\"\"\n    sequences: dict id -> sequence\n    ids: ordered list of ids\n    returns: np.ndarray [n_proteins, emb_dim]\n    \"\"\"\n    all_embs = []\n    \n    for start in range(0, len(ids), batch_size):\n        batch_ids = ids[start:start + batch_size]\n        batch_seqs = [sequences[i] for i in batch_ids]\n        \n        if MODEL_TYPE == 'esm2':\n            # Truncate to max length to avoid errors\n            batch_seqs_trunc = [s[:MAX_LENGTH] for s in batch_seqs]\n            \n            inputs = tokenizer(\n                batch_seqs_trunc,\n                return_tensors='pt',\n                padding=True,\n                truncation=True,\n                max_length=MAX_LENGTH,\n            ).to(device)\n            \n            outputs = model(**inputs)\n            # outputs.last_hidden_state: [batch, seq_len, hidden]\n            # Mean pool over non-padding tokens\n            attn_mask = inputs['attention_mask'].unsqueeze(-1)  # [batch, seq_len, 1]\n            summed = (outputs.last_hidden_state * attn_mask).sum(dim=1)\n            lengths = attn_mask.sum(dim=1)  # [batch, 1]\n            emb = (summed / lengths).cpu().numpy()  # [batch, hidden]\n            \n        elif MODEL_TYPE == 'esm3':\n            embeddings = get_cambrian_embeddings(batch_seqs)\n\n        \n        all_embs.append(emb)\n        \n        if (start // batch_size) % 10 == 0:\n            print(f'Processed {start + len(batch_ids)} / {len(ids)} proteins')\n\n    return np.concatenate(all_embs, axis=0)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:42.718928Z","iopub.execute_input":"2025-11-30T21:23:42.719220Z","iopub.status.idle":"2025-11-30T21:23:42.728524Z","shell.execute_reply.started":"2025-11-30T21:23:42.719190Z","shell.execute_reply":"2025-11-30T21:23:42.727701Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------\n# 5. Run for train & test\n# ------------------\ntrain_ids = list(train_sequences.keys())\ntest_ids  = list(test_sequences.keys())\n\nprint(\"Computing train embeddings...\")\ntrain_embeddings = compute_embeddings(train_sequences, train_ids)\n\nprint(\"Computing test embeddings...\")\ntest_embeddings = compute_embeddings(test_sequences, test_ids)\n\nprint(f'Embeddings shapes -> Train: {train_embeddings.shape}, Test: {test_embeddings.shape}')\n\n# ------------------\n# 6. Save\n# ------------------\nout_dir = '/kaggle/working'\nnp.save(os.path.join(out_dir, 'protein_embeddings_train.npy'), train_embeddings)\nnp.save(os.path.join(out_dir, 'protein_embeddings_test.npy'),  test_embeddings)\nnp.save(os.path.join(out_dir, 'train_ids.npy'), np.array(train_ids))\nnp.save(os.path.join(out_dir, 'test_ids.npy'),  np.array(test_ids))\n\nprint(f'Saved embeddings and ids to {out_dir}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:42.729342Z","iopub.execute_input":"2025-11-30T21:23:42.729565Z","iopub.status.idle":"2025-11-30T21:23:55.348842Z","shell.execute_reply.started":"2025-11-30T21:23:42.729542Z","shell.execute_reply":"2025-11-30T21:23:55.348174Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Some visualization","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\n\n# Convert to DataFrames\ndf_train = pd.DataFrame(train_embeddings)\ndf_train[\"protein_id\"] = train_ids\ndf_train[\"split\"] = \"train\"\n\ndf_test = pd.DataFrame(test_embeddings)\ndf_test[\"protein_id\"] = test_ids\ndf_test[\"split\"] = \"test\"\n\n# Combine into one\ndf_all = pd.concat([df_train, df_test], axis=0, ignore_index=True)\n\nprint(df_all.shape)\ndf_all.head()\n# All columns except protein_id + split\nembedding_cols = df_all.columns[:-2]   # assumes last 2 are protein_id, split\nX = df_all[embedding_cols].values\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:55.349528Z","iopub.execute_input":"2025-11-30T21:23:55.349724Z","iopub.status.idle":"2025-11-30T21:23:55.360637Z","shell.execute_reply.started":"2025-11-30T21:23:55.349707Z","shell.execute_reply":"2025-11-30T21:23:55.359780Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.manifold import TSNE\nimport matplotlib.pyplot as plt\n\ntsne = TSNE(\n    n_components=2,\n    perplexity=30,\n    metric='cosine',\n    init='pca',\n    learning_rate='auto',\n    random_state=0\n)\n\ncoords_tsne = tsne.fit_transform(X)\ndf_all[\"TSNE1\"] = coords_tsne[:,0]\ndf_all[\"TSNE2\"] = coords_tsne[:,1]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:55.361456Z","iopub.execute_input":"2025-11-30T21:23:55.361701Z","iopub.status.idle":"2025-11-30T21:23:56.560751Z","shell.execute_reply.started":"2025-11-30T21:23:55.361678Z","shell.execute_reply":"2025-11-30T21:23:56.560099Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.figure(figsize=(8,7))\n\nfor split, color in [(\"train\",\"tab:blue\"), (\"test\",\"tab:orange\")]:\n    mask = df_all[\"split\"] == split\n    plt.scatter(\n        df_all.loc[mask, \"TSNE1\"],\n        df_all.loc[mask, \"TSNE2\"],\n        s=6,\n        alpha=0.7,\n        label=split,\n        c=color\n    )\n\nplt.legend()\nplt.xlabel(\"t-SNE 1\")\nplt.ylabel(\"t-SNE 2\")\nplt.title(\"t-SNE of Protein Embeddings (colored by train/test)\")\nplt.tight_layout()\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-30T21:23:56.562998Z","iopub.execute_input":"2025-11-30T21:23:56.563494Z","iopub.status.idle":"2025-11-30T21:23:56.791923Z","shell.execute_reply.started":"2025-11-30T21:23:56.563471Z","shell.execute_reply":"2025-11-30T21:23:56.791175Z"}},"outputs":[],"execution_count":null}]}