{"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":31328,"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\nfrom transformers import BertModel, BertTokenizer\n\nBASE_PATH = \"/kaggle/input/competitions/cafa-5-protein-function-prediction/Train/\"\n\n\ndef get_protbert_embeddings_multigpu():\n    \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    \n    df[['EntryID']].to_csv(\"train_ids.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(\"protbert_embeddings_full.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()\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}