{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":41875,"databundleVersionId":5521661},{"sourceType":"kernelVersion","sourceId":294261447}],"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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install biopython","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn # 引入神經網路模組\nimport pandas as pd\nimport numpy as np\nfrom Bio import SeqIO\nfrom transformers import AutoTokenizer, EsmModel, T5EncoderModel, T5Tokenizer\nfrom tqdm import tqdm\nimport gc\n\nBASE_PATH = \"/kaggle/input/competitions/cafa-5-protein-function-prediction/Train/\"\n\ndef get_embeddings_multigpu():\n    fasta_sequences = list(SeqIO.parse(BASE_PATH + \"train_sequences.fasta\", \"fasta\"))\n    df = pd.DataFrame({\n        'EntryID': [seq.id for seq in fasta_sequences],\n        'Sequence': [str(seq.seq) for seq in fasta_sequences]\n    })\n    df[['EntryID']].to_csv(\"train_ids.csv\", index=False)\n    \n    # 【關鍵 1】因為有兩張卡，Batch Size 可以放心放大兩倍\n    batch_size = 16 \n\n    # ==========================================\n    # 第一階段：ESM-2 雙卡平行處理\n    # ==========================================\n    print(f\"啟動 ESM-2 模型... (偵測到 {torch.cuda.device_count()} 張 GPU)\")\n    esm_model_name = \"facebook/esm2_t33_650M_UR50D\"   # \"facebook/esm2_t36_3B_UR50D\"\n    esm_tok = AutoTokenizer.from_pretrained(esm_model_name)\n    \n    # 先將模型以 FP16 載入 CPU，再放進 DataParallel，最後丟進 cuda\n    esm_model = EsmModel.from_pretrained(esm_model_name, torch_dtype=torch.float16)\n    esm_model = nn.DataParallel(esm_model).cuda() # 【關鍵 2】啟動雙卡平行\n    esm_model.eval()\n    \n    esm_list = []\n    for i in tqdm(range(0, len(df), batch_size), desc=\"ESM-2 (Dual-GPU)\"):\n        batch_seqs = df['Sequence'].iloc[i:i+batch_size].tolist()\n        inputs = esm_tok(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024).to(\"cuda\")\n        \n        with torch.no_grad(): \n            outputs = esm_model(**inputs)\n            esm_list.append(outputs.last_hidden_state.mean(dim=1).cpu().float().numpy())\n            \n    np.save(\"esm_embeddings_full.npy\", np.concatenate(esm_list))\n    print(\"ESM-2 提取完成！\")\n    \n    # 🧹 清空雙卡快取\n    del esm_model, esm_tok, inputs, outputs, esm_list\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    # ==========================================\n    # 第二階段：ProtT5 雙卡平行處理\n    # ==========================================\n    print(\"啟動 ProtT5 模型...\")\n    t5_model_name = \"Rostlab/prot_t5_xl_uniref50\"\n    t5_tok = T5Tokenizer.from_pretrained(t5_model_name, do_lower_case=False)\n    \n    t5_model = T5EncoderModel.from_pretrained(t5_model_name, torch_dtype=torch.float16)\n    t5_model = nn.DataParallel(t5_model).cuda() # 【關鍵 2】啟動雙卡平行\n    t5_model.eval()\n\n    t5_list = []\n    for i in tqdm(range(0, len(df), batch_size), desc=\"ProtT5 (Dual-GPU)\"):\n        batch_seqs = df['Sequence'].iloc[i:i+batch_size].tolist()\n        batch_seqs = [\" \".join(list(s)) for s in batch_seqs]\n        inputs = t5_tok(batch_seqs, return_tensors=\"pt\", padding=True).to(\"cuda\")\n        \n        with torch.no_grad(): \n            outputs = t5_model(**inputs)\n            t5_list.append(outputs.last_hidden_state.mean(dim=1).cpu().float().numpy())\n            \n    np.save(\"t5_embeddings_full.npy\", np.concatenate(t5_list))\n    print(\"ProtT5 提取完成！\")\n\nget_embeddings_multigpu()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport numpy as np\nimport scipy.sparse as sp\nfrom tqdm import tqdm\nimport os\n\n# --- 1. 設定檔案路徑 ---\nBASE_DIR = \"/kaggle/input/notebooks/mafiosoquasar/ems-protbert-model-ensemble/outputs_labels/\"\nPOS_KEYS_PATH = os.path.join(BASE_DIR, \"goa_positive_keys.json\")\nNEG_KEYS_PATH = os.path.join(BASE_DIR, \"goa_negative_keys.json\")\nMAPS_PATH = os.path.join(BASE_DIR, \"maps.json\")\n\n# 輸出路徑 (存在 Kaggle 的 working directory)\nOUT_DIR = \"./\" \n\n# --- 2. 讀取原始資料與分類 Maps ---\nprint(\"正在讀取資料...\")\nwith open(POS_KEYS_PATH, \"r\") as f:\n    pos_keys = json.load(f)\nwith open(NEG_KEYS_PATH, \"r\") as f:\n    neg_keys = json.load(f)\n    \n# 讀取 maps.json (假設結構為 {\"BP\": [...], \"CC\": [...], \"MF\": [...]})\nwith open(MAPS_PATH, \"r\") as f:\n    go_maps = json.load(f)\n\n# 將標籤轉成 set，大幅提升後續的搜尋速度\ngo_sets = {\n    \"BP\": set(go_maps.get(\"BP\", [])),\n    \"CC\": set(go_maps.get(\"CC\", [])),\n    \"MF\": set(go_maps.get(\"MF\", []))\n}\n\n# --- 3. 建立全局的 Protein ID 對照表 (Row Mapping) ---\nprint(\"正在建立全局 Protein ID 索引...\")\nall_proteins = set()\nfor key in tqdm(pos_keys + neg_keys, desc=\"掃描所有蛋白質\"):\n    prot, _ = key.split('_')\n    all_proteins.add(prot)\n\nproteins_list = sorted(list(all_proteins))\n# ★ 這就是你之前苦苦尋找的 Protein ID Mapping！\nprotein_to_idx = {prot: i for i, prot in enumerate(proteins_list)}\nprint(f\"✅ 全局蛋白質總數: {len(proteins_list)}\")\n\n# 將 Protein Mapping 獨立存下來\nwith open(os.path.join(OUT_DIR, \"protein_mapping.json\"), \"w\") as f:\n    json.dump(protein_to_idx, f)\n\n\n# --- 4. 針對 BP, CC, MF 分別建立矩陣 ---\ndef build_and_save_ontology_matrix(ontology_name, go_set):\n    if not go_set:\n        print(f\"⚠️ {ontology_name} 沒有找到 GO 標籤，跳過。\")\n        return\n\n    print(f\"\\n[{ontology_name}] 開始處理...\")\n    \n    # 1. 整理該類別專屬的 GO 標籤清單 (Column Mapping)\n    go_list = sorted(list(go_set))\n    go_to_idx = {go: i for i, go in enumerate(go_list)}\n    \n    matrix_shape = (len(proteins_list), len(go_list))\n    print(f\"[{ontology_name}] 矩陣大小預計為: {matrix_shape} (Proteins x GO_Terms)\")\n    \n    # 2. 過濾並建立矩陣的內部函數\n    def create_sparse(keys_list, val):\n        row_idx, col_idx, data = [], [], []\n        for key in keys_list:\n            prot, go = key.split('_')\n            # 只有當這個 GO 標籤屬於當前分類 (BP/CC/MF) 時才加入\n            if go in go_set: \n                row_idx.append(protein_to_idx[prot])\n                col_idx.append(go_to_idx[go])\n                data.append(val)\n        return sp.coo_matrix((data, (row_idx, col_idx)), shape=matrix_shape).tocsr()\n\n    # 3. 建立正負樣本矩陣\n    pos_matrix = create_sparse(tqdm(pos_keys, desc=f\"過濾 {ontology_name} 正樣本\"), val=1)\n    neg_matrix = create_sparse(tqdm(neg_keys, desc=f\"過濾 {ontology_name} 負樣本\"), val=1)\n    \n    # 4. 儲存 NPZ 檔案 (檔名與你要求的一致)\n    pos_file = os.path.join(OUT_DIR, f\"labels_{ontology_name}.npz\")\n    sp.save_npz(pos_file, pos_matrix)\n    print(f\"[{ontology_name}] ✅ 已儲存: {pos_file} (包含 {pos_matrix.nnz} 個正標籤)\")\n    \n    sp.save_npz(os.path.join(OUT_DIR, f\"labels_{ontology_name}_neg.npz\"), neg_matrix)\n\n# --- 5. 執行三大本體論的處理 ---\nfor ont in [\"BP\", \"CC\", \"MF\"]:\n    build_and_save_ontology_matrix(ont, go_sets[ont])\n\nprint(\"全部任務完成！你現在擁有三個乾淨的 NPZ 矩陣，以及對應的 Protein ID Mapping 了！\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}