{"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}],"dockerImageVersionId":31329,"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\n!pip install obonet","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/Test (Targets)/\"\n\ndef get_embeddings_multigpu():\n    fasta_sequences = list(SeqIO.parse(BASE_PATH + \"testsuperset.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(\"test_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(\"test_esm_emb_aligned.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\nget_embeddings_multigpu()","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\nfrom transformers import BertModel, BertTokenizer\n\nBASE_PATH = \"/kaggle/input/competitions/cafa-5-protein-function-prediction/Test (Targets)/\"\n\n\ndef get_protbert_embeddings_multigpu():\n    \n    fasta_sequences = list(SeqIO.parse(BASE_PATH + \"testsuperset.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    \n    df[['EntryID']].to_csv(\"test_ids_2.csv\", index=False)\n    \n    batch_size = 16 \n\n    # ==========================================\n    # 第二階段：ProtBERT 雙卡平行處理\n    # ==========================================\n    print(f\"🚀 啟動 ProtBERT 模型... (共 {len(df)} 筆序列)\")\n    bert_model_name = \"Rostlab/prot_bert\"\n    \n    # 換成 BertTokenizer\n    bert_tok = BertTokenizer.from_pretrained(bert_model_name, do_lower_case=False)\n    \n    # 換成 BertModel (保持 float16 運算加速且省顯存)\n    bert_model = BertModel.from_pretrained(bert_model_name, torch_dtype=torch.float16)\n    bert_model = nn.DataParallel(bert_model).cuda() # 【關鍵 2】啟動雙卡平行\n    bert_model.eval()\n\n    bert_list = []\n    for i in tqdm(range(0, len(df), batch_size), desc=\"ProtBERT (Dual-GPU)\"):\n        batch_seqs = df['Sequence'].iloc[i:i+batch_size].tolist()\n        \n        # ProtBERT 同樣要求胺基酸之間要有空白\n        batch_seqs = [\" \".join(list(s)) for s in batch_seqs]\n        \n        inputs = bert_tok(batch_seqs, return_tensors=\"pt\", padding=True).to(\"cuda\")\n        inputs = bert_tok(\n                    batch_seqs, \n                    return_tensors=\"pt\", \n                    padding=True, \n                    truncation=True,    \n                    max_length=2048     \n                ).to(\"cuda\")\n        \n        with torch.no_grad(): \n            outputs = bert_model(**inputs)\n            # 取 Mean Pooling (ProtBERT 預設輸出也是 1024 維)\n            features = outputs.last_hidden_state.mean(dim=1).cpu().float().numpy()\n            bert_list.append(features)\n            \n        # 偶爾清理一下記憶體碎片\n        if i % (batch_size * 100) == 0:\n            torch.cuda.empty_cache()\n            \n    final_emb = np.concatenate(bert_list)\n    \n    # 【關鍵 3】安全檢查：確保 Embedding 總數跟 ID 總數一致才存檔\n    if final_emb.shape[0] == len(df):\n        np.save(\"test_pb_emb_aligned.npy\", final_emb)\n        print(f\"✅ ProtBERT 提取完成！最終形狀: {final_emb.shape} (維度應為 1024)\")\n    else:\n        print(f\"❌ 警告：Embedding 數量 ({final_emb.shape[0]}) 與 ID 數量 ({len(df)}) 不符！\")\n\n# 執行\nget_protbert_embeddings_multigpu()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nid_1 = pd.read_csv(\"test_ids.csv\")\nid_2 = pd.read_csv(\"test_ids_2.csv\")\n\n\nif id_1.equals(id_2):\n    print(\"✅ 對齊完成！\")\nelse:\n    print(\"❌ 對齊錯誤！\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}