{"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":87793,"databundleVersionId":12276181,"sourceType":"competition"},{"sourceId":11797207,"sourceType":"datasetVersion","datasetId":7408208}],"dockerImageVersionId":31040,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Import libraries and set up paths\nimport os\nimport sys\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport random\nfrom tqdm.notebook import tqdm\nimport yaml\nimport matplotlib.pyplot as plt\n\n# Set random seeds for reproducibility\ntorch.manual_seed(0)\nnp.random.seed(0)\nrandom.seed(0)\n\n# Define paths\nCOMPETITION_PATH = \"/kaggle/input/stanford-rna-3d-folding\"\nWORKSPACE_PATH = \"/kaggle/input/ribonanzanet3d/ribonanzanet3D\"\nMODEL_DIR = f\"{WORKSPACE_PATH}/models\"\nOUTPUT_DIR = f\"{WORKSPACE_PATH}/output\"\n\n# Add model path to Python path\nsys.path.append(f\"{MODEL_DIR}/ribonanzanet2d-final\")\n\n# Show directory structure\nprint(\"Working directory contents:\")\n!ls -la {WORKSPACE_PATH}\nprint(\"\\nModel directory contents:\")\n!ls -la {MODEL_DIR}\nprint(\"\\nOutput directory contents:\")\n!ls -la {OUTPUT_DIR}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T21:16:43.842583Z","iopub.execute_input":"2025-05-13T21:16:43.842997Z","iopub.status.idle":"2025-05-13T21:16:44.252739Z","shell.execute_reply.started":"2025-05-13T21:16:43.842896Z","shell.execute_reply":"2025-05-13T21:16:44.251724Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define model architecture\nclass Config:\n    def __init__(self, **entries):\n        self.__dict__.update(entries)\n        self.entries = entries\n\ndef load_config_from_yaml(file_path):\n    with open(file_path, 'r') as file:\n        config = yaml.safe_load(file)\n    return Config(**config)\n\n# Check if Network.py exists\nif not os.path.exists(f\"{MODEL_DIR}/ribonanzanet2d-final/Network.py\"):\n    print(\"Network.py not found. Please create it using the %%writefile magic command or upload it.\")\nelse:\n    try:\n        from Network import RibonanzaNet\n        print(\"Successfully imported RibonanzaNet!\")\n        \n        class finetuned_RibonanzaNet(RibonanzaNet):\n            def __init__(self, config, pretrained=False):\n                config.dropout = 0.1\n                super(finetuned_RibonanzaNet, self).__init__(config)\n                self.dropout = nn.Dropout(0.0)\n                self.xyz_predictor = nn.Linear(256, 3)\n            \n            def forward(self, src):\n                sequence_features, pairwise_features = self.get_embeddings(\n                    src, torch.ones_like(src).long().to(src.device))\n                xyz = self.xyz_predictor(sequence_features)\n                return xyz\n    except Exception as e:\n        print(f\"Error importing RibonanzaNet: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T21:16:49.675067Z","iopub.execute_input":"2025-05-13T21:16:49.675398Z","iopub.status.idle":"2025-05-13T21:16:49.684173Z","shell.execute_reply.started":"2025-05-13T21:16:49.675375Z","shell.execute_reply":"2025-05-13T21:16:49.682973Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load test data\ntest_data = pd.read_csv(f\"{COMPETITION_PATH}/test_sequences.csv\")\nprint(f\"Loaded {len(test_data)} test sequences\")\n\n# Define dataset class\nclass RNADataset(Dataset):\n    def __init__(self, data):\n        self.data = data\n        self.tokens = {nt:i for i,nt in enumerate('ACGU')}\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        sequence = [self.tokens[nt] for nt in (self.data.loc[idx, 'sequence'])]\n        sequence = np.array(sequence)\n        sequence = torch.tensor(sequence)\n        return {'sequence': sequence}\n\n# Create test dataset\ntest_dataset = RNADataset(test_data)\nprint(f\"Example sequence length: {len(test_dataset[0]['sequence'])}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T21:16:52.378431Z","iopub.execute_input":"2025-05-13T21:16:52.378734Z","iopub.status.idle":"2025-05-13T21:16:52.398692Z","shell.execute_reply.started":"2025-05-13T21:16:52.378709Z","shell.execute_reply":"2025-05-13T21:16:52.397728Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load model\n# Check if config file exists\nconfig_file = f\"{MODEL_DIR}/ribonanzanet2d-final/configs/pairwise.yaml\"\nif not os.path.exists(config_file):\n    print(f\"Config file not found at {config_file}\")\n    # Create a basic config if needed\n    basic_config = {\n        \"n_tokens\": 4,\n        \"d_model\": 256,\n        \"d_ff\": 1024,\n        \"n_layers\": 9,\n        \"n_heads\": 8,\n        \"dropout\": 0.1,\n        \"max_len\": 384\n    }\n    os.makedirs(os.path.dirname(config_file), exist_ok=True)\n    with open(config_file, 'w') as f:\n        yaml.dump(basic_config, f)\n    print(f\"Created basic config file at {config_file}\")\n\n# Check if model weights exist\nmodel_file = f\"{OUTPUT_DIR}/RibonanzaNet-3D.pt\"\nif not os.path.exists(model_file):\n    print(f\"Model file not found at {model_file}\")\n    print(\"Please upload the model weights file\")\nelse:\n    print(f\"Found model weights at {model_file}\")\n\ntry:\n    # Initialize model\n    config = load_config_from_yaml(config_file)\n    model = finetuned_RibonanzaNet(config, pretrained=False)\n    \n    # Check if CUDA is available\n    if torch.cuda.is_available():\n        model = model.cuda()\n        print(\"Using GPU for inference\")\n    else:\n        print(\"Using CPU for inference\")\n    \n    # Load pre-trained weights\n    model.load_state_dict(torch.load(model_file, \n                                    map_location=\"cuda\" if torch.cuda.is_available() else \"cpu\"))\n    print(\"Model loaded successfully!\")\n\n    # Generate predictions\n    model.eval()\n    preds = []\n    print(\"Generating predictions...\")\n\n    for i in tqdm(range(len(test_dataset))):\n        src = test_dataset[i]['sequence'].long()\n        src = src.unsqueeze(0)\n        if torch.cuda.is_available():\n            src = src.cuda()\n        \n        # Generate 5 predictions per sequence\n        tmp = []\n        # First 4 with dropout enabled (stochastic)\n        model.train()\n        for j in range(4):\n            with torch.no_grad():\n                xyz = model(src).squeeze()\n            tmp.append(xyz.cpu().numpy())\n        \n        # Last one without dropout\n        model.eval()\n        with torch.no_grad():\n            xyz = model(src).squeeze()\n        tmp.append(xyz.cpu().numpy())\n        \n        tmp = np.stack(tmp, 0)\n        preds.append(tmp)\n\n    print(f\"Generated predictions for {len(preds)} sequences\")\n\n    # Format submission\n    print(\"Formatting submission...\")\n    data = []\n    \n    for i in range(len(test_data)):\n        for j in range(len(test_data.loc[i, 'sequence'])):\n            row = [\n                test_data.loc[i, 'target_id'] + f\"_{j+1}\",\n                test_data.loc[i, 'sequence'][j],\n                j+1  # 1-indexed\n            ]\n            \n            for k in range(5):  # 5 predictions\n                for coord in range(3):  # x, y, z\n                    row.append(preds[i][k][j][coord])\n            \n            data.append(row)\n    \n    columns = ['ID', 'resname', 'resid']\n    for i in range(1, 6):\n        columns += [f\"x_{i}\", f\"y_{i}\", f\"z_{i}\"]\n    \n    submission = pd.DataFrame(data, columns=columns)\n    submission.to_csv('submission.csv', index=False)\n    print(\"Submission saved!\")\n    \n    # Display first few rows\n    submission.head()\n    \nexcept Exception as e:\n    print(f\"Error in model loading or inference: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T21:16:54.487348Z","iopub.execute_input":"2025-05-13T21:16:54.488131Z","iopub.status.idle":"2025-05-13T21:16:54.676827Z","shell.execute_reply.started":"2025-05-13T21:16:54.488093Z","shell.execute_reply":"2025-05-13T21:16:54.675954Z"}},"outputs":[],"execution_count":null}]}