{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\n\n# 输入输出路径\nfile_path = \"/kaggle/input/leash-BELKA/train.csv\"\noutput_file = \"binds_1.csv\"\n\nCHUNK_SIZE = 100_0000  # 每次读取的行数，根据内存调整\nfirst_write = True    # 控制是否写入表头\nMAX_CHUNKS = 20      # 最多处理 200 个 chunk\n\nprint(\"开始提取 binds=1 的行...\")\n\nwith pd.read_csv(file_path, chunksize=CHUNK_SIZE) as reader:\n    for i, chunk in enumerate(reader):\n        print(f\"正在处理第 {i+1} 块（共 {CHUNK_SIZE} 行）...\")\n        \n        # 筛选 binds=1 的行\n        filtered_chunk = chunk[chunk['binds'] == 1]\n\n        # 写入文件，第一块写入时包含表头\n        filtered_chunk.to_csv(output_file, mode='w' if first_write else 'a',\n                              index=False, header=first_write)\n        first_write = False\n        \n        # 达到最大块数后终止\n        if i + 1 >= MAX_CHUNKS:\n            print(f\"已达到最大块数 {MAX_CHUNKS}，提前终止。\")\n            break\n\nprint(f\"提取完成，所有 binds=1 的行已保存至 '{output_file}'\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:02:37.621049Z","iopub.execute_input":"2025-06-03T05:02:37.621484Z","iopub.status.idle":"2025-06-03T05:04:05.390921Z","shell.execute_reply.started":"2025-06-03T05:02:37.621456Z","shell.execute_reply":"2025-06-03T05:04:05.389964Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\n# 文件路径\npos_file = \"binds_1.csv\"\ntrain_file = \"/kaggle/input/leash-BELKA/train.csv\"\noutput_file = \"balanced_train.csv\"\n\nCHUNK_SIZE = 100_00\n\n# 设置采样数量\nnum_pos_needed = 10_000\nprint(f\"设定抽取的正样本数量 (binds=1): {num_pos_needed:,}\")\n\n# 1️⃣ 从 binds_1.csv 中随机读取 50 万个正样本\nprint(\"开始从 binds_1.csv 中读取并抽样正样本...\")\npos_df = pd.read_csv(pos_file)\nif len(pos_df) < num_pos_needed:\n    print(f\"⚠️ 注意：binds_1.csv 中只有 {len(pos_df):,} 条正样本，不足 {num_pos_needed:,}\")\n    num_pos_needed = len(pos_df)\n\npos_sampled = pos_df.sample(n=num_pos_needed, random_state=42)\n\n# 2️⃣ 从 train.csv 中筛选 binds=0 的行，并随机抽取 50 万条\nneg_collector = []\ntotal_neg_collected = 0\ntarget_neg = 9*num_pos_needed  # 和正样本数量一致\n\nprint(f\"开始从 train.csv 中读取并抽样 {target_neg:,} 条负样本...\")\n\nwith pd.read_csv(train_file, chunksize=CHUNK_SIZE, iterator=True) as reader:\n    for i, chunk in enumerate(reader):\n        print(f\"处理第 {i+1} 块...\")\n        neg_chunk = chunk[chunk['binds'] == 0]\n        if not neg_chunk.empty:\n            neg_collector.append(neg_chunk)\n            total_neg_collected += len(neg_chunk)\n\n            if total_neg_collected >= target_neg:\n                print(\"已收集足够数量的负样本，停止读取...\")\n                break\n\n# 合并所有负样本块\nif neg_collector:\n    neg_df = pd.concat(neg_collector)\n    if len(neg_df) > target_neg:\n        neg_df = neg_df.sample(n=target_neg, random_state=42)\n    else:\n        print(f\"⚠️ 注意：只找到 {len(neg_df):,} 条负样本，不足 {target_neg:,}\")\nelse:\n    raise ValueError(\"未找到任何负样本，请检查 train.csv 中是否存在 binds=0 的行。\")\n\n# 3️⃣ 合并正负样本并打乱顺序\nprint(\"合并正样本和负样本，并打乱顺序...\")\nfinal_df = pd.concat([pos_sampled, neg_df], ignore_index=True).sample(frac=1.0, random_state=42)\n\n# 4️⃣ 保存结果\nprint(f\"正在保存平衡后的数据集至 '{output_file}'\")\nfinal_df.to_csv(output_file, index=False)\n\nprint(f\"✅ 平衡完成！共 {len(final_df):,} 行数据，其中：\")\nprint(f\"   - 正样本: {len(pos_sampled):,}\")\nprint(f\"   - 负样本: {len(neg_df):,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:05.392596Z","iopub.execute_input":"2025-06-03T05:04:05.392988Z","iopub.status.idle":"2025-06-03T05:04:06.809329Z","shell.execute_reply.started":"2025-06-03T05:04:05.392955Z","shell.execute_reply":"2025-06-03T05:04:06.808185Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -*- coding: utf-8 -*-\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, confusion_matrix\nimport math\nimport re","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:06.810633Z","iopub.execute_input":"2025-06-03T05:04:06.810917Z","iopub.status.idle":"2025-06-03T05:04:06.832631Z","shell.execute_reply.started":"2025-06-03T05:04:06.810893Z","shell.execute_reply":"2025-06-03T05:04:06.831434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 定义 SMILES Tokenizer\ndef smi_tokenizer(smi):\n    pattern = \"(\\$$|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\\||$|$|\\.|=|#|-|\\+|\\\\|\\/|:|~|@|\\?|>>?|\\*|\\$|[0-9])\"\n    tokens = [token for token in re.split(pattern, smi) if token != '']\n    return tokens","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:06.834537Z","iopub.execute_input":"2025-06-03T05:04:06.834831Z","iopub.status.idle":"2025-06-03T05:04:06.840451Z","shell.execute_reply.started":"2025-06-03T05:04:06.834808Z","shell.execute_reply":"2025-06-03T05:04:06.839348Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 加载数据\ndf = pd.read_csv(\"balanced_train.csv\")\nprint(df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:06.841425Z","iopub.execute_input":"2025-06-03T05:04:06.841734Z","iopub.status.idle":"2025-06-03T05:04:07.115868Z","shell.execute_reply.started":"2025-06-03T05:04:06.841709Z","shell.execute_reply":"2025-06-03T05:04:07.114879Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 划分训练集和验证集\ntrain_df, val_df = train_test_split(\n    df,\n    test_size=0.3,\n    random_state=42,\n    stratify=df['binds']\n)\n# 构建词汇表\nall_smiles = df['molecule_smiles'].unique()\nvocab = set()\nfor smi in all_smiles:\n    vocab.update(smi_tokenizer(smi))\nvocab = ['<PAD>', '<UNK>', '<CLS>'] + sorted(vocab)\ntoken2idx = {t: i for i, t in enumerate(vocab)}\nVOCAB_SIZE = len(token2idx)\nMAX_LEN = 128\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:07.116871Z","iopub.execute_input":"2025-06-03T05:04:07.117134Z","iopub.status.idle":"2025-06-03T05:04:08.019988Z","shell.execute_reply.started":"2025-06-03T05:04:07.117114Z","shell.execute_reply":"2025-06-03T05:04:08.018941Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------------\n# Step 3: 自定义 Dataset\n# -------------------------------\nclass MoleculeDataset(Dataset):\n    def __init__(self, dataframe, max_len=MAX_LEN, tokenizer=smi_tokenizer, token2idx=token2idx):\n        self.df = dataframe.reset_index(drop=True)\n        self.max_len = max_len\n        self.tokenizer = tokenizer\n        self.token2idx = token2idx\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        smi = row['molecule_smiles']\n        label = row['binds']\n\n        tokens = ['<CLS>'] + self.tokenizer(smi)\n        tokens = tokens[:self.max_len]\n        input_ids = [self.token2idx.get(t, self.token2idx['<UNK>']) for t in tokens]\n        input_ids += [self.token2idx['<PAD>']] * (self.max_len - len(input_ids))\n        attention_mask = [1] * len(tokens) + [0] * (self.max_len - len(tokens))\n\n        return (\n            torch.tensor(input_ids, dtype=torch.long),\n            torch.tensor(attention_mask, dtype=torch.long),\n            torch.tensor(label, dtype=torch.float)\n        )\n\n# 创建 DataLoader\ntrain_dataset = MoleculeDataset(train_df)\nval_dataset = MoleculeDataset(val_df)\n\ntrain_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=0, pin_memory=True)\nval_loader = DataLoader(val_dataset, batch_size=128, num_workers=0, pin_memory=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:08.020912Z","iopub.execute_input":"2025-06-03T05:04:08.021161Z","iopub.status.idle":"2025-06-03T05:04:08.035273Z","shell.execute_reply.started":"2025-06-03T05:04:08.021143Z","shell.execute_reply":"2025-06-03T05:04:08.034434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------------\n# Step 4: 定义 PositionalEncoding\n# -------------------------------\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=512):\n        super().__init__()\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        x = x + self.pe[:x.size(1)]\n        return x\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:08.036067Z","iopub.execute_input":"2025-06-03T05:04:08.036288Z","iopub.status.idle":"2025-06-03T05:04:08.058487Z","shell.execute_reply.started":"2025-06-03T05:04:08.036271Z","shell.execute_reply":"2025-06-03T05:04:08.057296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------------\n# Step 5: 定义 Transformer 模型\n# -------------------------------\nclass TransformerModel(nn.Module):\n    def __init__(self, vocab_size, embed_dim=128, num_heads=4, num_layers=2, dim_feedforward=256):\n        super().__init__()\n        self.embedding = nn.Embedding(vocab_size, embed_dim)\n        self.positional_encoding = PositionalEncoding(embed_dim)\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=embed_dim,\n            nhead=num_heads,\n            dim_feedforward=dim_feedforward\n        )\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n        self.classifier = nn.Linear(embed_dim, 1)\n\n    def forward(self, src, src_key_padding_mask=None):\n        src = self.embedding(src) * math.sqrt(self.embedding.embedding_dim)\n        src = self.positional_encoding(src)\n        src = src.permute(1, 0, 2)  # [seq_len, batch_size, embed_dim]\n        output = self.transformer_encoder(src, src_key_padding_mask=src_key_padding_mask)\n        output = output.mean(dim=0)  # 或者取 [CLS] token\n        return self.classifier(output)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:08.059346Z","iopub.execute_input":"2025-06-03T05:04:08.059620Z","iopub.status.idle":"2025-06-03T05:04:08.084445Z","shell.execute_reply.started":"2025-06-03T05:04:08.059599Z","shell.execute_reply":"2025-06-03T05:04:08.083167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------------\n# Step 6: 设置 device 和损失函数\n# -------------------------------\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\n# 计算类别权重\ny = df['binds'].values\nclass_weight = len(y[y == 0]) / len(y[y == 1])\nweight = torch.tensor([class_weight], dtype=torch.float32).to(device)\n\n# 定义损失函数\ncriterion = nn.BCEWithLogitsLoss(pos_weight=weight)\n\n# 初始化模型\nmodel = TransformerModel(VOCAB_SIZE).to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:08.087269Z","iopub.execute_input":"2025-06-03T05:04:08.088549Z","iopub.status.idle":"2025-06-03T05:04:11.782048Z","shell.execute_reply.started":"2025-06-03T05:04:08.088468Z","shell.execute_reply":"2025-06-03T05:04:11.776078Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nimport numpy as np\nfrom sklearn.metrics import roc_auc_score, accuracy_score\nimport torch\n\ndef train_model(model, train_loader, val_loader, criterion, optimizer, device, num_epochs=10, patience=5, save_path=\"best_model.pth\"):\n    model.to(device)\n    best_auc = 0\n    counter = 0\n    history = {\n        \"train_loss\": [],\n        \"val_loss\": [],\n        \"val_auc\": [],\n        \"val_acc\": []\n    }\n\n    for epoch in range(num_epochs):\n        # 每个 epoch 创建一个新的进度条\n        pbar = tqdm(total=len(train_loader), \n                    desc=f\"Epoch {epoch+1}/{num_epochs}\", \n                    position=0, \n                    leave=True)\n\n        # ----------------------\n        # 训练阶段\n        # ----------------------\n        model.train()\n        running_loss = 0.0\n\n        for inputs, masks, labels in train_loader:\n            inputs = inputs.to(device)\n            masks = masks.to(device)\n            labels = labels.to(device)\n\n            optimizer.zero_grad()\n\n            outputs = model(inputs, src_key_padding_mask=(masks == 0)).squeeze()\n            loss = criterion(outputs, labels)\n\n            loss.backward()\n            optimizer.step()\n\n            running_loss += loss.item() * inputs.size(0)\n\n            # 更新进度条\n            pbar.update(1)\n\n        epoch_loss = running_loss / len(train_loader.dataset)\n        history[\"train_loss\"].append(epoch_loss)\n\n        # ----------------------\n        # 验证阶段\n        # ----------------------\n        model.eval()\n        all_preds, all_labels = [], []\n        val_loss = 0.0\n\n        with torch.no_grad():\n            for inputs, masks, labels in val_loader:\n                inputs = inputs.to(device)\n                masks = masks.to(device)\n                labels = labels.to(device)\n\n                outputs = model(inputs, src_key_padding_mask=(masks == 0)).squeeze()\n                loss = criterion(outputs, labels)\n\n                val_loss += loss.item() * inputs.size(0)\n\n                preds = torch.sigmoid(outputs).cpu().numpy()\n                all_preds.extend(preds)\n                all_labels.extend(labels.cpu().numpy())\n\n        val_loss_avg = val_loss / len(val_loader.dataset)\n        history[\"val_loss\"].append(val_loss_avg)\n\n        # 计算 AUC 和 Accuracy\n        all_preds = np.array(all_preds)\n        all_labels = np.array(all_labels)\n        all_preds_class = (all_preds > 0.5).astype(int)\n        auc = roc_auc_score(all_labels, all_preds)\n        acc = accuracy_score(all_labels, all_preds_class)\n\n        history[\"val_auc\"].append(auc)\n        history[\"val_acc\"].append(acc)\n\n        # 完成当前 epoch 后关闭进度条，并打印最终结果\n        pbar.set_description(f\"Epoch {epoch+1}/{num_epochs} | Train Loss: {epoch_loss:.4f} | Val Loss: {val_loss_avg:.4f} | AUC: {auc:.4f} | Acc: {acc:.4f}\")\n        pbar.close()\n\n        # ----------------------\n        # 模型保存与早停机制\n        # ----------------------\n        if auc > best_auc:\n            best_auc = auc\n            torch.save(model.state_dict(), save_path)\n            counter = 0\n        else:\n            counter += 1\n\n        if counter >= patience:\n            print(\"Early stopping triggered.\")\n            break\n\n    print(f\"Training finished. Best AUC: {best_auc:.4f}\")\n    return history","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:11.782900Z","iopub.execute_input":"2025-06-03T05:04:11.783502Z","iopub.status.idle":"2025-06-03T05:04:11.800006Z","shell.execute_reply.started":"2025-06-03T05:04:11.783473Z","shell.execute_reply":"2025-06-03T05:04:11.798570Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = TransformerModel(VOCAB_SIZE).to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.BCEWithLogitsLoss(pos_weight=weight)\n\nhistory = train_model(\n    model,\n    train_loader=train_loader,\n    val_loader=val_loader,\n    criterion=criterion,\n    optimizer=optimizer,\n    device=device,\n    num_epochs=10,\n    patience=5,\n    save_path=\"best_model.pth\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T05:04:11.801406Z","iopub.execute_input":"2025-06-03T05:04:11.801963Z","iopub.status.idle":"2025-06-03T06:08:02.348808Z","shell.execute_reply.started":"2025-06-03T05:04:11.801919Z","shell.execute_reply":"2025-06-03T06:08:02.346093Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"test","metadata":{}},{"cell_type":"code","source":"import torch\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\n\n# 假设你已经在前面定义了 MoleculeDataset 和 TransformerModel 类\n# 如果没有，请复制粘贴它们到这个脚本中\n\n# -------------------------------\n# Step 1: 加载测试数据\n# -------------------------------\ntest_df = pd.read_csv(\"/kaggle/input/leash-BELKA/test.csv\")\n\n# 添加 binds 列用于 Dataset 兼容（不需要真实值）\ntest_df['binds'] = -1\n\n# 创建 Dataset 和 DataLoader\ntest_dataset = MoleculeDataset(test_df)\ntest_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)\n\n# -------------------------------\n# Step 2: 加载模型\n# -------------------------------\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = TransformerModel(VOCAB_SIZE).to(device)\nmodel.load_state_dict(torch.load(\"best_model.pth\"))  # 加载最佳模型\nmodel.eval()\n\n# -------------------------------\n# Step 3: 进行预测\n# -------------------------------\nall_ids = []\nall_preds = []\n\nwith torch.no_grad():\n    for inputs, masks, _ in tqdm(test_loader, desc=\"Predicting\"):\n        inputs = inputs.to(device)\n        masks = masks.to(device)\n\n        outputs = model(inputs, src_key_padding_mask=(masks == 0)).squeeze()\n        probs = torch.sigmoid(outputs).cpu().numpy()  # 概率输出\n        # classes = (probs > 0.5).astype(int)  # 如果你需要二分类结果，取消注释这行\n\n        all_preds.extend(probs)\n        all_ids.extend(test_df.iloc[:len(inputs)].id.values)\n\n# -------------------------------\n# Step 4: 构建结果并保存\n# -------------------------------\nresults_df = pd.DataFrame({\n    'id': all_ids,\n    'bind': all_preds  # 或者使用 classes 输出 0/1\n})\n\n# 保存为 CSV\nresults_df.to_csv(\"submission.csv\", index=False)\n\nprint(\"预测完成，结果已保存至 submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T07:01:59.202326Z","iopub.execute_input":"2025-06-03T07:01:59.202733Z","iopub.status.idle":"2025-06-03T07:52:24.944440Z","shell.execute_reply.started":"2025-06-03T07:01:59.202707Z","shell.execute_reply":"2025-06-03T07:52:24.943349Z"}},"outputs":[],"execution_count":null}]}