{"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":67356,"databundleVersionId":8006601,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":847738,"databundleVersionId":16911899,"modelInstanceId":644612,"modelId":656568,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":855716,"databundleVersionId":17025107,"modelInstanceId":650335,"modelId":662340,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":856203,"databundleVersionId":17032609,"modelInstanceId":650706,"modelId":662702,"isSourceIdPinned":false},{"sourceType":"kernelVersion","sourceId":313811183,"isSourceIdPinned":false}],"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 transformers torch rdkit","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport torch\nimport polars as pl\nimport pandas as pd\nfrom tqdm import tqdm\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.amp import autocast\nfrom transformers import DataCollatorWithPadding\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nfrom tqdm import tqdm\nfrom transformers import AutoTokenizer, AutoModel\nfrom sklearn.metrics import average_precision_score\nfrom sklearn.model_selection import train_test_split\nimport pyarrow.parquet as pq\nfrom torch.amp import GradScaler","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 1. 配置與蛋白質預快取 ---\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nPROT_MODEL_NAME = \"facebook/esm2_t12_35M_UR50D\"\n\n# 預先載入 Tokenizer\nprot_tokenizer = AutoTokenizer.from_pretrained(PROT_MODEL_NAME)\nprot_model = AutoModel.from_pretrained(PROT_MODEL_NAME).to(device).eval()\n\ndef precompute_prot_embeddings():\n    \"\"\"預先計算 BELKA 的三種蛋白質 Embedding\"\"\"\n    prot_sequences = {\n        'BRD4': \"MSAESGPGTRLRNLPVMGDGLETSQMSTTQAQAQPQPANAASTNPPPPETSNPNKPKRQTNQLQYLLRVVLKTLWKHQFAWPFQQPVDAVKLNLPDYYKIIKTPMDMGTIKKRLENNYYWNAQECIQDFNTMFTNCYIYNKPGDDIVLMAEALEKLFLQKINELPTEETEIMIVQAKGRGRGRKETGTAKPGVSTVPNTTQASTPPQTQTPQPNPPPVQATPHPFPAVTPDLIVQTPVMTVVPPQPLQTPPPVPPQPQPPPAPAPQPVQSHPPIIAATPQPVKTKKGVKRKADTTTPTTIDPIHEPPSLPPEPKTTKLGQRRESSRPVKPPKKDVPDSQQHPAPEKSSKVSEQLKCCSGILKEMFAKKHAAYAWPFYKPVDVEALGLHDYCDIIKHPMDMSTIKSKLEAREYRDAQEFGADVRLMFSNCYKYNPPDHEVVAMARKLQDVFEMRFAKMPDEPEEPVVAVSSPAVPPPTKVVAPPSSSDSSSDSSSDSDSSTDDSEEERAQRLAELQEQLKAVHEQLAALSQPQQNKPKKKKKKKKKKKK\", \n        'HSA': \"MKWVTFISLLFLFSSAYSRGVFRRDAHKSEVAHRFKDLGEENFKALVLIAFAQYLQQCPFEDHVKLVNEVTEFAKTCVADESAENCDKSLHTLFGDKLCTVATLRETYGEMADCCAKQEPERNECFLQHKDDNPNLPRLVRPEVDVMCTAFHDNEETFLKKYLYEIARRHPYFYAPELLFFAKRYKAAFTECCQAADKAACLLPKLDELRDEGKASSAKQRLKCASLQKFGERAFKAWAVARLSQRFPKAEFAEVSKLVTDLTKVHTECCHGDLLECADDRADLAKYICENQDSISSKLKECCEKPLLEKSHCIAEVENDEMPADLPSLAADFVESKDVCKNYAEAKDVFLGMFLYEYARRHPDYSVVLLLRLAKTYETTLEKCCAAADPHECYAKVFDEFKPLVEEPQNLIKQNCELFEQLGEYKFQNALLVRYTKKVPQVSTPTLVEVSRNLGKVGSKCCKHPEAKRMPCAEDYLSVVLNQLCVLHEKTPVSDRVTKCCTESLVNRRPCFSALEVDETYVPKEFNAETFTFHADICTLSEKERQIKKQTALVELVKHKPKATKEQLKAVMDDFAAFVEKCCKADDKETCFAEEGKKLVAASQAALGL\",\n        'sEH': \"MTLRAAVFDLDGVLALPAVFGVLGRTEEALALPRGLLNDAFQKGGPEGATTRLMKGEITLSQWIPLMEENCRKCSETAKVCLPKNFSIKEIFDKAISARKINRPMLQAALMLRKKGFTTAILTNTWLDDRAERDGLAQLMCELKMHFDFLIESCQVGMVKPEPQIYKFLLDTLKASPSEVVFLDDIGANLKPARDLGMVTILVQDTDTALKELEKVTGIQLLNTPAPLPTSCNPSDMSHGYVTVKPRVRLHFVELGSGPAVCLCHGFPESWYSWRYQIPALAQAGYRVLAMDMKGYGESSAPPEIEEYCMEVLCKEMVTFLDKLGLSQAVFIGHDWGGMLVWYMALFYPERVRAVASLNTPFIPANPNMSPLESIKANPVFDYQLYFQEPGVAEAELEQNLSRTFKSLFRASDESVLSMHKVCEAGGLFVNSPEEPSLSRMVTEEEIQFYVQQFKKSGFRGPLNWYRNMERNWKWACKSLGRKILIPALMVTAEKDFVLVPQMSQHMEDWIPHLKRGHIEDCGHWTQMDKPTEVNQILIKWLDSDARNPPVVSKM\"\n    }\n    prot_map = {}\n    for name, seq in prot_sequences.items():\n        inputs = prot_tokenizer(seq, return_tensors=\"pt\", truncation=True, max_length=1024).to(device)\n        with torch.no_grad():\n            outputs = prot_model(**inputs)\n            # 取 Mean Pooling\n            prot_map[name] = outputs.last_hidden_state.mean(dim=1).cpu().numpy().squeeze()\n    return prot_map\n\n# 預處理好的 map\nprot_map = precompute_prot_embeddings()\n\nclass SmilesCharTokenizer:\n    def __init__(self):\n        # BELKA 常見的 SMILES 字元庫 (包含我們保留的 Dy 標籤)\n        chars = ['<PAD>', 'C', 'N', 'O', 'F', 'S', 'c', 'n', 'o', 's', \n                 '(', ')', '[', ']', '=', '#', '@', '+', '-', '1', '2', '3', \n                 '4', '5', '6', '7', '8', '9', 'l', 'r', 'B', 'D', 'y', 'H']\n        \n        self.vocab = {c: i for i, c in enumerate(chars)}\n        self.vocab_size = len(self.vocab)\n        self.pad_token_id = 0\n\n    def __call__(self, smiles_list, max_length=128):\n        batch_ids = []\n        for smiles in smiles_list:\n            # 將字串轉換為 ID，遇到沒見過的字元統一當作 padding 或是忽略\n            token_ids = [self.vocab.get(c, 0) for c in smiles]\n            \n            # 截斷 (Truncation)\n            if len(token_ids) > max_length:\n                token_ids = token_ids[:max_length]\n            \n            batch_ids.append(token_ids)\n            \n        # 動態 Padding (Padding 到該 Batch 的最大長度)\n        max_len_in_batch = max(len(t) for t in batch_ids)\n        padded_ids = []\n        for t in batch_ids:\n            padded = t + [self.pad_token_id] * (max_len_in_batch - len(t))\n            padded_ids.append(padded)\n            \n        return {\"input_ids\": torch.tensor(padded_ids, dtype=torch.long)}\n\n# 實例化新的 Tokenizer\nmol_tokenizer = SmilesCharTokenizer()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BELKAHybridModel(nn.Module):\n    def __init__(self, vocab_size, embed_dim=128, cnn_dim=128, prot_dim=480, max_len=256):\n        super().__init__()\n        \n        # 1. 基礎字元 Embedding\n        # padding_idx=0 確保 <PAD> token 的初始向量為全 0\n        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)\n        \n        # 2. 局部特徵提取器：1D-CNN (Kernel=3)\n        # 模仿第五名：提取相鄰 3 個字元 (例如 'C', '=', 'O' 形成羰基) 的組合特徵\n        self.conv1d = nn.Conv1d(in_channels=embed_dim, out_channels=cnn_dim, kernel_size=3, padding=1)\n        self.cnn_activation = nn.GELU()\n        \n        # 3. 空間位置編碼 (Positional Encoding)\n        # 雖然 CNN 有局部順序，但 Transformer 看整個序列是無序的，必須加上位置編碼\n        self.pos_encoder = nn.Embedding(max_len, cnn_dim)\n        \n        # 4. 全局特徵編碼器：極度輕量的 Transformer Encoder (4 層, 4 頭)\n        # 模仿第一名：讓局部官能基互相產生 Attention，理解 2D 分子拓撲\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=cnn_dim, \n            nhead=4, \n            dim_feedforward=cnn_dim * 4, \n            batch_first=True, \n            dropout=0.1,\n            activation='gelu'\n        )\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=4)\n        \n        # 5. 終極預測頭 (分類器)\n        self.head = nn.Sequential(\n            nn.Linear(cnn_dim + prot_dim, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 128),\n            nn.BatchNorm1d(128),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(128, 1)\n        )\n\n    def forward(self, mol_ids, prot_emb):\n        \"\"\"\n        mol_ids: (Batch, Seq_Len) - 字元 token ID\n        prot_emb: (Batch, Prot_Dim) - 蛋白質特徵向量\n        \"\"\"\n        B, L = mol_ids.size()\n        \n        # --- Stage 1: Embedding ---\n        x = self.embedding(mol_ids)  # Shape: (B, L, Embed_Dim)\n        \n        # --- Stage 2: 1D-CNN 特徵提取 ---\n        # PyTorch 的 Conv1d 要求輸入維度為 (Batch, Channels, Length)\n        x = x.permute(0, 2, 1)       # 轉換為 (B, Embed_Dim, L)\n        x = self.conv1d(x)           # 卷積提取 -> (B, CNN_Dim, L)\n        x = self.cnn_activation(x)\n        x = x.permute(0, 2, 1)       # 轉回 Transformer 需要的 (B, L, CNN_Dim)\n        \n        # --- Stage 3: 加入位置編碼 ---\n        positions = torch.arange(L, device=x.device).unsqueeze(0).expand(B, L)\n        x = x + self.pos_encoder(positions)\n        \n        # --- Stage 4: Transformer 處理 ---\n        # 創建 Padding Mask：mol_ids == 0 的地方為 True，告訴 Transformer「不要看這裡」\n        pad_mask = (mol_ids == 0)    \n        \n        # 輸入 Transformer，src_key_padding_mask 負責擋掉無效的空白區域\n        x = self.transformer(x, src_key_padding_mask=pad_mask) # Shape: (B, L, CNN_Dim)\n        \n        # --- Stage 5: Masked Average Pooling (遮罩平均池化) ---\n        # 絕對不能直接 .mean(dim=1)！因為後面一大堆 pad token 是 0，會把平均值稀釋掉！\n        # 必須只計算「有效字元」的平均。\n        token_mask = (~pad_mask).unsqueeze(-1).float()          # 變成 1 (有效) 和 0 (無效), Shape: (B, L, 1)\n        sum_x = (x * token_mask).sum(dim=1)                     # 將有效字元的特徵相加, Shape: (B, CNN_Dim)\n        sum_mask = token_mask.sum(dim=1).clamp(min=1e-9)        # 計算有效字元的總數，clamp 防止除以 0\n        pooled_x = sum_x / sum_mask                             # 真正的有效平均特徵！Shape: (B, CNN_Dim)\n        \n        # --- Stage 6: 結合蛋白質特徵並預測 ---\n        combined = torch.cat([pooled_x, prot_emb], dim=1)       # Shape: (B, CNN_Dim + Prot_Dim)\n        logits = self.head(combined)                            # Shape: (B, 1)\n        \n        return logits","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport torch\nimport polars as pl\nimport pandas as pd\nfrom tqdm import tqdm\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.amp import autocast\nfrom transformers import DataCollatorWithPadding\n\n\nclass BELKATestDataset(Dataset):\n    def __init__(self, df, prot_map):\n        self.ids = df['id'].to_numpy()\n        self.smiles = df['molecule_smiles'].to_numpy()\n        self.prot_names = df['protein_name'].to_numpy()\n        self.prot_map = prot_map\n\n    def __len__(self):\n        return len(self.smiles)\n\n    def __getitem__(self, idx):\n        # 1. 抓取蛋白質特徵\n        prot_vector = self.prot_map[self.prot_names[idx]]\n        if isinstance(prot_vector, torch.Tensor):\n            prot_vector = prot_vector.clone().detach().float()\n        else:\n            prot_vector = torch.tensor(prot_vector, dtype=torch.float32)\n\n        return {\n            \"id\": self.ids[idx],\n            \"smiles\": self.smiles[idx], \n            \"prot_emb\": prot_vector\n        }\n\n# ==========================================\n# 動態 Tokenize 的 Collate Function\n# ==========================================\ndef create_test_collate_fn(tokenizer):\n    def collate_fn(batch):\n        ids = [item['id'] for item in batch]\n        smiles_list = [item['smiles'] for item in batch]\n        prot_embs = torch.stack([item['prot_emb'] for item in batch])\n        \n        tokenized = tokenizer(smiles_list, max_length=256)\n        mol_ids = tokenized[\"input_ids\"]\n        \n        return {\n            \"ids\": ids, \n            \"mol_ids\": mol_ids,\n            \"prot_emb\": prot_embs\n        }\n    return collate_fn\n\n# ==========================================\n# generate_submission\n# ==========================================\ndef generate_submission_optimized(model, test_file, mol_tokenizer, prot_map, device, output_name=\"submission.csv\"):\n    model.eval()\n    \n    print(\"使用 Polars 讀取測試集中...\")\n    test_df = pl.read_parquet(test_file, columns=[\"id\", \"molecule_smiles\", \"protein_name\"])\n    \n    print(\"準備 DataLoader...\")\n    test_dataset = BELKATestDataset(test_df, prot_map)\n    test_collate = create_test_collate_fn(mol_tokenizer)\n    \n    # 測試集的 Batch Size 可以開很大，因為不需要計算梯度 (Backprop)\n    test_loader = DataLoader(\n        test_dataset, \n        batch_size=4096, \n        shuffle=False, # 不洗牌！\n        num_workers=4, \n        pin_memory=True,\n        collate_fn=test_collate\n    )\n    \n    all_ids = []\n    all_preds = []\n\n    print(\"開始進行推論 (Inference)...\")\n    with torch.no_grad():\n        # 使用 torch.amp.autocast 加速推論\n        for batch in tqdm(test_loader, desc=\"Predicting\"):\n            batch_ids = batch[\"ids\"]\n            mol_ids = batch[\"mol_ids\"].to(device)\n            prot_emb = batch[\"prot_emb\"].to(device)\n            \n            with torch.amp.autocast('cuda'):\n                outputs = model(mol_ids=mol_ids, prot_emb=prot_emb)\n            \n            # 將 logits 轉成 0~1 的機率\n            probs = torch.sigmoid(outputs).cpu().numpy().flatten()\n            \n            all_ids.extend(batch_ids)\n            all_preds.extend(probs)\n\n            \n    print(f\"正在寫入 {output_name}...\")\n    # 使用 Pandas 快速寫出 CSV\n    sub_df = pd.DataFrame({\n        \"id\": all_ids,\n        \"binds\": all_preds\n    })\n    sub_df.to_csv(output_name, index=False)\n    print(\"✅ 提交檔案生成完畢！祝你 Leaderboard 暴衝！\")\n\n    del test_df, all_ids, all_preds, sub_df\n    gc.collect()\n\n# ==========================================\n# 執行區塊\n# ==========================================\n\nmodel = BELKAHybridModel(\n    vocab_size=mol_tokenizer.vocab_size, \n    embed_dim=128, \n    cnn_dim=256, \n    prot_dim=480\n).to(device)\n\nmodel.load_state_dict(torch.load(\"/kaggle/input/models/t8101349/model-v4/pytorch/default/1/best_model_chemberta_1_4.pth\", map_location=device))\n\nTEST_FILE = \"/kaggle/input/competitions/leash-BELKA/test.parquet\" \nSAMPLE_SUB_FILE = \"/kaggle/input/competitions/leash-BELKA/sample_submission.csv\"\n\nif os.path.exists(TEST_FILE):\n    generate_submission_optimized(model, TEST_FILE, mol_tokenizer, prot_map, device)\nelse:\n    print(f\"❌ 找不到測試檔: {TEST_FILE}，請檢查路徑。\")\n\n# 最後檢查行數\nif os.path.exists(\"submission.csv\") and os.path.exists(SAMPLE_SUB_FILE):\n    # 用 Polars 掃描行數最快，不吃記憶體\n    sub_count = pl.scan_csv(\"submission.csv\").select(pl.len()).collect().item()\n    sample_count = pl.scan_csv(SAMPLE_SUB_FILE).select(pl.len()).collect().item()\n    \n    if sub_count == sample_count:\n        print(f\"✅ 行數檢查完美通過: {sub_count} 筆數據\")\n    else:\n        print(f\"❌ 警告：行數不符！提交檔: {sub_count}, 範本: {sample_count}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import polars as pl\nsub = pl.read_csv(\"submission.csv\")\nprint(\"ID 是否嚴格遞增且無遺漏：\", sub['id'].is_sorted())\n\n# 把 submission 的結果跟 test.parquet 裡的 protein_name 對接起來看平均機率\ntest_df = pl.read_parquet(\"/kaggle/input/competitions/leash-BELKA/test.parquet\", columns=[\"id\", \"protein_name\"])\nsub_merged = sub.join(test_df, on=\"id\")\nprint(sub_merged.group_by(\"protein_name\").agg(pl.col(\"binds\").mean()))","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}