{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false}],"dockerImageVersionId":31286,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\n# ==========================================\n# 1. Config & Hyperparameters\n# ==========================================\nclass Config:\n    seed = 42\n    max_len = 512       # Adjust based on GPU memory and max RNA sequence length\n    vocab_size = 6      # A, C, G, U, N (Unknown), PAD\n    embed_dim = 256\n    num_heads = 8\n    num_layers = 6\n    dropout = 0.1\n    batch_size = 4\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ncfg = Config()\ntorch.manual_seed(cfg.seed)\n\n# RNA vocabulary mapping\nVOCAB = {'A': 1, 'C': 2, 'G': 3, 'U': 4, 'N': 5}\nPAD_IDX = 0\n\ndef seq_to_tensor(seq, max_len):\n    \"\"\"Convert RNA sequence string to an integer Tensor\"\"\"\n    tokens = [VOCAB.get(s, 5) for s in seq]\n    if len(tokens) < max_len:\n        tokens += [PAD_IDX] * (max_len - len(tokens))\n    else:\n        tokens = tokens[:max_len]\n    return torch.tensor(tokens, dtype=torch.long)\n\n# ==========================================\n# 2. Deep Learning Model (Transformer)\n# ==========================================\nclass RNA3DPredictor(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.embedding = nn.Embedding(cfg.vocab_size, cfg.embed_dim, padding_idx=PAD_IDX)\n        self.pos_encoder = nn.Parameter(torch.randn(1, cfg.max_len, cfg.embed_dim))\n        \n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=cfg.embed_dim, \n            nhead=cfg.num_heads, \n            dim_feedforward=cfg.embed_dim * 4, \n            dropout=cfg.dropout, \n            batch_first=True\n        )\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=cfg.num_layers)\n        \n        # Output layer: predict 5 structures * 3 coordinates (x,y,z) = 15 values per residue\n        self.regressor = nn.Sequential(\n            nn.Linear(cfg.embed_dim, cfg.embed_dim // 2),\n            nn.GELU(),\n            nn.Linear(cfg.embed_dim // 2, 5 * 3) \n        )\n\n    def forward(self, x, mask=None):\n        # x shape: (Batch, Seq_Len)\n        seq_len = x.size(1)\n        emb = self.embedding(x) + self.pos_encoder[:, :seq_len, :]\n        \n        # Transformer processing\n        out = self.transformer(emb, src_key_padding_mask=mask)\n        \n        # Coordinate prediction shape: (Batch, Seq_Len, 15)\n        coords = self.regressor(out)\n        return coords\n\n# ==========================================\n# 3. Inference & Submission Pipeline (Fixed for Long Sequences)\n# ==========================================\ndef generate_submission(model, test_csv_path, sample_sub_path, output_path):\n    model.eval()\n    \n    # Read sample submission to get required IDs and residue info\n    sub_df = pd.read_csv(sample_sub_path)\n    \n    # Extract Target ID and Residue ID\n    sub_df['target_id'] = sub_df['ID'].apply(lambda x: x.rsplit('_', 1)[0])\n    sub_df['resid'] = sub_df['ID'].apply(lambda x: int(x.rsplit('_', 1)[1]))\n    \n    predictions = []\n    \n    # Group by target_id to process each RNA sequence\n    grouped = sub_df.groupby('target_id')\n    \n    with torch.no_grad():\n        for target_id, group in grouped:\n            group = group.sort_values('resid')\n            seq = \"\".join(group['resname'].values)\n            \n            # Convert to Tensor (this will truncate to cfg.max_len if too long)\n            x = seq_to_tensor(seq, cfg.max_len).unsqueeze(0).to(cfg.device)\n            \n            # Model prediction: shape (1, min(Seq_Len, Max_Len), 15)\n            preds = model(x).squeeze(0).cpu().numpy()\n            \n            # Extract valid predictions\n            valid_preds = preds\n            \n            for i, (_, row) in enumerate(group.iterrows()):\n                pred_row = {\n                    'ID': row['ID'],\n                    'resname': row['resname'],\n                    'resid': row['resid']\n                }\n                \n                coord_idx = 0\n                for copy_idx in range(1, 6):\n                    # 【核心修复】检查当前残基是否超出了我们预测的最大长度\n                    if i < valid_preds.shape[0]:\n                        pred_row[f'x_{copy_idx}'] = valid_preds[i, coord_idx]\n                        pred_row[f'y_{copy_idx}'] = valid_preds[i, coord_idx + 1]\n                        pred_row[f'z_{copy_idx}'] = valid_preds[i, coord_idx + 2]\n                    else:\n                        # 如果超出，填充 0.0 作为安全后备，防止代码崩溃\n                        pred_row[f'x_{copy_idx}'] = 0.0\n                        pred_row[f'y_{copy_idx}'] = 0.0\n                        pred_row[f'z_{copy_idx}'] = 0.0\n                    coord_idx += 3\n                \n                predictions.append(pred_row)\n\n    # Save final results\n    pred_df = pd.DataFrame(predictions)\n    \n    # Ensure column order matches exactly\n    cols = ['ID', 'resname', 'resid', \n            'x_1', 'y_1', 'z_1', 'x_2', 'y_2', 'z_2', \n            'x_3', 'y_3', 'z_3', 'x_4', 'y_4', 'z_4', 'x_5', 'y_5', 'z_5']\n    pred_df = pred_df[cols]\n    \n    pred_df.to_csv(output_path, index=False)\n    print(f\"✅ Submission saved to {output_path}\")\n\n# ==========================================\n# 4. Main Execution\n# ==========================================\nif __name__ == \"__main__\":\n    # Initialize model and load to GPU\n    model = RNA3DPredictor(cfg).to(cfg.device)\n    \n    # Bulletproof function to automatically find files in Kaggle's input directory\n    def get_kaggle_path(filename, search_dir=\"/kaggle/input\"):\n        for root, dirs, files in os.walk(search_dir):\n            if filename in files:\n                return os.path.join(root, filename)\n        return filename # Fallback to current directory if not found\n    \n    # Automatically locate the exact paths\n    TEST_SEQ_PATH = get_kaggle_path(\"test_sequences.csv\")\n    SAMPLE_SUB_PATH = get_kaggle_path(\"sample_submission.csv\")\n    OUTPUT_PATH = \"submission.csv\"\n    \n    # Print the paths to verify where Kaggle actually stored them\n    print(f\"🔍 Found Test Sequences at: {TEST_SEQ_PATH}\")\n    print(f\"🔍 Found Sample Submission at: {SAMPLE_SUB_PATH}\")\n    \n    generate_submission(model, TEST_SEQ_PATH, SAMPLE_SUB_PATH, OUTPUT_PATH)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-20T20:28:17.237473Z","iopub.execute_input":"2026-03-20T20:28:17.237799Z","iopub.status.idle":"2026-03-20T20:28:21.973902Z","shell.execute_reply.started":"2026-03-20T20:28:17.237771Z","shell.execute_reply":"2026-03-20T20:28:21.972746Z"}},"outputs":[],"execution_count":null}]}