{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12024591,"sourceType":"competition"},{"sourceId":5123458,"sourceType":"datasetVersion","datasetId":2975803},{"sourceId":10880419,"sourceType":"datasetVersion","datasetId":6760509},{"sourceId":10923077,"sourceType":"datasetVersion","datasetId":6785143},{"sourceId":10924604,"sourceType":"datasetVersion","datasetId":6791615}],"dockerImageVersionId":30919,"isInternetEnabled":false,"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,"execution":{"iopub.status.busy":"2025-03-07T15:37:47.306761Z","iopub.execute_input":"2025-03-07T15:37:47.307012Z","iopub.status.idle":"2025-03-07T15:37:48.065897Z","shell.execute_reply.started":"2025-03-07T15:37:47.306978Z","shell.execute_reply":"2025-03-07T15:37:48.065065Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Environments\n","metadata":{}},{"cell_type":"code","source":"%ls /kaggle/input/boltz-dependencies","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:37:48.066745Z","iopub.execute_input":"2025-03-07T15:37:48.067147Z","iopub.status.idle":"2025-03-07T15:37:48.189685Z","shell.execute_reply.started":"2025-03-07T15:37:48.067123Z","shell.execute_reply":"2025-03-07T15:37:48.188652Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/boltz-dependencies/*whl --no-deps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:37:48.190944Z","iopub.execute_input":"2025-03-07T15:37:48.191300Z","iopub.status.idle":"2025-03-07T15:37:57.935381Z","shell.execute_reply.started":"2025-03-07T15:37:48.191268Z","shell.execute_reply":"2025-03-07T15:37:57.934517Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/fairscale-0413/*whl --no-deps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:37:57.937845Z","iopub.execute_input":"2025-03-07T15:37:57.938085Z","iopub.status.idle":"2025-03-07T15:37:59.183318Z","shell.execute_reply.started":"2025-03-07T15:37:57.938063Z","shell.execute_reply":"2025-03-07T15:37:59.182501Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/biopython/*whl --no-deps","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prepare scripts","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/working/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:37:59.184852Z","iopub.execute_input":"2025-03-07T15:37:59.185130Z","iopub.status.idle":"2025-03-07T15:37:59.190376Z","shell.execute_reply.started":"2025-03-07T15:37:59.185108Z","shell.execute_reply":"2025-03-07T15:37:59.189588Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%mkdir inputs_prediction\n%mkdir outputs_prediction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:37:59.191155Z","iopub.execute_input":"2025-03-07T15:37:59.191368Z","iopub.status.idle":"2025-03-07T15:37:59.434690Z","shell.execute_reply.started":"2025-03-07T15:37:59.191349Z","shell.execute_reply":"2025-03-07T15:37:59.433763Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cp -rf /kaggle/input/rna-prediction-boltz/boltz/src/boltz .","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:37:59.435529Z","iopub.execute_input":"2025-03-07T15:37:59.435789Z","iopub.status.idle":"2025-03-07T15:38:00.429364Z","shell.execute_reply.started":"2025-03-07T15:37:59.435766Z","shell.execute_reply":"2025-03-07T15:38:00.428313Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%ls boltz","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:38:00.430382Z","iopub.execute_input":"2025-03-07T15:38:00.430718Z","iopub.status.idle":"2025-03-07T15:38:00.551020Z","shell.execute_reply.started":"2025-03-07T15:38:00.430686Z","shell.execute_reply":"2025-03-07T15:38:00.549829Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Write file","metadata":{}},{"cell_type":"code","source":"%%writefile inference.py\n\nimport pickle\nimport urllib.request\nfrom dataclasses import asdict, dataclass\nfrom pathlib import Path\nfrom typing import Literal, Optional\n\nimport click\nimport torch\nfrom pytorch_lightning import Trainer, seed_everything\nfrom pytorch_lightning.strategies import DDPStrategy\nfrom pytorch_lightning.utilities import rank_zero_only\nfrom tqdm import tqdm\n\nfrom boltz.data import const\nfrom boltz.data.module.inference import BoltzInferenceDataModule\nfrom boltz.data.msa.mmseqs2 import run_mmseqs2\nfrom boltz.data.parse.a3m import parse_a3m\nfrom boltz.data.parse.csv import parse_csv\nfrom boltz.data.parse.fasta import parse_fasta\nfrom boltz.data.parse.yaml import parse_yaml\nfrom boltz.data.types import MSA, Manifest, Record\nfrom boltz.data.write.writer import BoltzWriter\nfrom boltz.model.model import Boltz1\n\nCCD_URL = \"https://huggingface.co/boltz-community/boltz-1/resolve/main/ccd.pkl\"\nMODEL_URL = (\n    \"https://huggingface.co/boltz-community/boltz-1/resolve/main/boltz1_conf.ckpt\"\n)\n\n\n@dataclass\nclass BoltzProcessedInput:\n    \"\"\"Processed input data.\"\"\"\n\n    manifest: Manifest\n    targets_dir: Path\n    msa_dir: Path\n\n\n@dataclass\nclass BoltzDiffusionParams:\n    \"\"\"Diffusion process parameters.\"\"\"\n\n    gamma_0: float = 0.605\n    gamma_min: float = 1.107\n    noise_scale: float = 0.901\n    rho: float = 8\n    step_scale: float = 1.638\n    sigma_min: float = 0.0004\n    sigma_max: float = 160.0\n    sigma_data: float = 16.0\n    P_mean: float = -1.2\n    P_std: float = 1.5\n    coordinate_augmentation: bool = True\n    alignment_reverse_diff: bool = True\n    synchronize_sigmas: bool = True\n    use_inference_model_cache: bool = True\n\n\n@rank_zero_only\ndef download(cache: Path) -> None:\n    \"\"\"Download all the required data.\n\n    Parameters\n    ----------\n    cache : Path\n        The cache directory.\n\n    \"\"\"\n    # Download CCD\n    ccd = cache / \"ccd.pkl\"\n    if not ccd.exists():\n        click.echo(\n            f\"Downloading the CCD dictionary to {ccd}. You may \"\n            \"change the cache directory with the --cache flag.\"\n        )\n        urllib.request.urlretrieve(CCD_URL, str(ccd))  # noqa: S310\n\n    # Download model\n    model = cache / \"boltz1_conf.ckpt\"\n    if not model.exists():\n        click.echo(\n            f\"Downloading the model weights to {model}. You may \"\n            \"change the cache directory with the --cache flag.\"\n        )\n        urllib.request.urlretrieve(MODEL_URL, str(model))  # noqa: S310\n\n\ndef check_inputs(\n    data: Path,\n    outdir: Path,\n    override: bool = False,\n) -> list[Path]:\n    \"\"\"Check the input data and output directory.\n\n    If the input data is a directory, it will be expanded\n    to all files in this directory. Then, we check if there\n    are any existing predictions and remove them from the\n    list of input data, unless the override flag is set.\n\n    Parameters\n    ----------\n    data : Path\n        The input data.\n    outdir : Path\n        The output directory.\n    override: bool\n        Whether to override existing predictions.\n\n    Returns\n    -------\n    list[Path]\n        The list of input data.\n\n    \"\"\"\n    click.echo(\"Checking input data.\")\n\n    # Check if data is a directory\n    if data.is_dir():\n        data: list[Path] = list(data.glob(\"*\"))\n\n        # Filter out non .fasta or .yaml files, raise\n        # an error on directory and other file types\n        filtered_data = []\n        for d in data:\n            if d.suffix in (\".fa\", \".fas\", \".fasta\", \".yml\", \".yaml\"):\n                filtered_data.append(d)\n            elif d.is_dir():\n                msg = f\"Found directory {d} instead of .fasta or .yaml.\"\n                raise RuntimeError(msg)\n            else:\n                msg = (\n                    f\"Unable to parse filetype {d.suffix}, \"\n                    \"please provide a .fasta or .yaml file.\"\n                )\n                raise RuntimeError(msg)\n\n        data = filtered_data\n    else:\n        data = [data]\n\n    # Check if existing predictions are found\n    existing = (outdir / \"predictions\").rglob(\"*\")\n    existing = {e.name for e in existing if e.is_dir()}\n\n    # Remove them from the input data\n    if existing and not override:\n        data = [d for d in data if d.stem not in existing]\n        num_skipped = len(existing) - len(data)\n        msg = (\n            f\"Found some existing predictions ({num_skipped}), \"\n            f\"skipping and running only the missing ones, \"\n            \"if any. If you wish to override these existing \"\n            \"predictions, please set the --override flag.\"\n        )\n        click.echo(msg)\n    elif existing and override:\n        msg = \"Found existing predictions, will override.\"\n        click.echo(msg)\n\n    return data\n\n\ndef compute_msa(\n    data: dict[str, str],\n    target_id: str,\n    msa_dir: Path,\n    msa_server_url: str,\n    msa_pairing_strategy: str,\n) -> None:\n    \"\"\"Compute the MSA for the input data.\n\n    Parameters\n    ----------\n    data : dict[str, str]\n        The input protein sequences.\n    target_id : str\n        The target id.\n    msa_dir : Path\n        The msa directory.\n    msa_server_url : str\n        The MSA server URL.\n    msa_pairing_strategy : str\n        The MSA pairing strategy.\n\n    \"\"\"\n    if len(data) > 1:\n        paired_msas = run_mmseqs2(\n            list(data.values()),\n            msa_dir / f\"{target_id}_paired_tmp\",\n            use_env=True,\n            use_pairing=True,\n            host_url=msa_server_url,\n            pairing_strategy=msa_pairing_strategy,\n        )\n    else:\n        paired_msas = [\"\"] * len(data)\n\n    unpaired_msa = run_mmseqs2(\n        list(data.values()),\n        msa_dir / f\"{target_id}_unpaired_tmp\",\n        use_env=True,\n        use_pairing=False,\n        host_url=msa_server_url,\n        pairing_strategy=msa_pairing_strategy,\n    )\n\n    for idx, name in enumerate(data):\n        # Get paired sequences\n        paired = paired_msas[idx].strip().splitlines()\n        paired = paired[1::2]  # ignore headers\n        paired = paired[: const.max_paired_seqs]\n\n        # Set key per row and remove empty sequences\n        keys = [idx for idx, s in enumerate(paired) if s != \"-\" * len(s)]\n        paired = [s for s in paired if s != \"-\" * len(s)]\n\n        # Combine paired-unpaired sequences\n        unpaired = unpaired_msa[idx].strip().splitlines()\n        unpaired = unpaired[1::2]\n        unpaired = unpaired[: (const.max_msa_seqs - len(paired))]\n        if paired:\n            unpaired = unpaired[1:]  # ignore query is already present\n\n        # Combine\n        seqs = paired + unpaired\n        keys = keys + [-1] * len(unpaired)\n\n        # Dump MSA\n        csv_str = [\"key,sequence\"] + [f\"{key},{seq}\" for key, seq in zip(keys, seqs)]\n\n        msa_path = msa_dir / f\"{name}.csv\"\n        with msa_path.open(\"w\") as f:\n            f.write(\"\\n\".join(csv_str))\n\n\n@rank_zero_only\ndef process_inputs(  # noqa: C901, PLR0912, PLR0915\n    data: list[Path],\n    out_dir: Path,\n    ccd_path: Path,\n    msa_server_url: str,\n    msa_pairing_strategy: str,\n    max_msa_seqs: int = 4096,\n    use_msa_server: bool = False,\n) -> None:\n    \"\"\"Process the input data and output directory.\n\n    Parameters\n    ----------\n    data : list[Path]\n        The input data.\n    out_dir : Path\n        The output directory.\n    ccd_path : Path\n        The path to the CCD dictionary.\n    max_msa_seqs : int, optional\n        Max number of MSA sequences, by default 4096.\n    use_msa_server : bool, optional\n        Whether to use the MMSeqs2 server for MSA generation, by default False.\n\n    Returns\n    -------\n    BoltzProcessedInput\n        The processed input data.\n\n    \"\"\"\n    click.echo(\"Processing input data.\")\n    existing_records = None\n\n    # Check if manifest exists at output path\n    manifest_path = out_dir / \"processed\" / \"manifest.json\"\n    if manifest_path.exists():\n        click.echo(f\"Found a manifest file at output directory: {out_dir}\")\n\n        manifest: Manifest = Manifest.load(manifest_path)\n        input_ids = [d.stem for d in data]\n        existing_records, processed_ids = zip(\n            *[\n                (record, record.id)\n                for record in manifest.records\n                if record.id in input_ids\n            ]\n        )\n\n        if isinstance(existing_records, tuple):\n            existing_records = list(existing_records)\n\n        # Check how many examples need to be processed\n        missing = len(input_ids) - len(processed_ids)\n        if not missing:\n            click.echo(\"All examples in data are processed. Updating the manifest\")\n            # Dump updated manifest\n            updated_manifest = Manifest(existing_records)\n            updated_manifest.dump(out_dir / \"processed\" / \"manifest.json\")\n            return\n\n        click.echo(f\"{missing} missing ids. Preprocessing these ids\")\n        missing_ids = list(set(input_ids).difference(set(processed_ids)))\n        data = [d for d in data if d.stem in missing_ids]\n        assert len(data) == len(missing_ids)\n\n    # Create output directories\n    msa_dir = out_dir / \"msa\"\n    structure_dir = out_dir / \"processed\" / \"structures\"\n    processed_msa_dir = out_dir / \"processed\" / \"msa\"\n    predictions_dir = out_dir / \"predictions\"\n\n    out_dir.mkdir(parents=True, exist_ok=True)\n    msa_dir.mkdir(parents=True, exist_ok=True)\n    structure_dir.mkdir(parents=True, exist_ok=True)\n    processed_msa_dir.mkdir(parents=True, exist_ok=True)\n    predictions_dir.mkdir(parents=True, exist_ok=True)\n\n    # Load CCD\n    with ccd_path.open(\"rb\") as file:\n        ccd = pickle.load(file)  # noqa: S301\n\n    if existing_records is not None:\n        click.echo(f\"Found {len(existing_records)} records. Adding them to records\")\n\n    # Parse input data\n    records: list[Record] = existing_records if existing_records is not None else []\n    for path in tqdm(data):\n        try:\n            # Parse data\n            if path.suffix in (\".fa\", \".fas\", \".fasta\"):\n                target = parse_fasta(path, ccd)\n            elif path.suffix in (\".yml\", \".yaml\"):\n                target = parse_yaml(path, ccd)\n            elif path.is_dir():\n                msg = f\"Found directory {path} instead of .fasta or .yaml, skipping.\"\n                raise RuntimeError(msg)\n            else:\n                msg = (\n                    f\"Unable to parse filetype {path.suffix}, \"\n                    \"please provide a .fasta or .yaml file.\"\n                )\n                raise RuntimeError(msg)\n\n            # Get target id\n            target_id = target.record.id\n\n            # Get all MSA ids and decide whether to generate MSA\n            to_generate = {}\n            prot_id = const.chain_type_ids[\"PROTEIN\"]\n            for chain in target.record.chains:\n                # Add to generate list, assigning entity id\n                if (chain.mol_type == prot_id) and (chain.msa_id == 0):\n                    entity_id = chain.entity_id\n                    msa_id = f\"{target_id}_{entity_id}\"\n                    to_generate[msa_id] = target.sequences[entity_id]\n                    chain.msa_id = msa_dir / f\"{msa_id}.csv\"\n\n                # We do not support msa generation for non-protein chains\n                elif chain.msa_id == 0:\n                    chain.msa_id = -1\n\n            # Generate MSA\n            if to_generate and not use_msa_server:\n                msg = \"Missing MSA's in input and --use_msa_server flag not set.\"\n                raise RuntimeError(msg)\n\n            if to_generate:\n                msg = f\"Generating MSA for {path} with {len(to_generate)} protein entities.\"\n                click.echo(msg)\n                compute_msa(\n                    data=to_generate,\n                    target_id=target_id,\n                    msa_dir=msa_dir,\n                    msa_server_url=msa_server_url,\n                    msa_pairing_strategy=msa_pairing_strategy,\n                )\n\n            # Parse MSA data\n            msas = sorted({c.msa_id for c in target.record.chains if c.msa_id != -1})\n            msa_id_map = {}\n            for msa_idx, msa_id in enumerate(msas):\n                # Check that raw MSA exists\n                msa_path = Path(msa_id)\n                if not msa_path.exists():\n                    msg = f\"MSA file {msa_path} not found.\"\n                    raise FileNotFoundError(msg)\n\n                # Dump processed MSA\n                processed = processed_msa_dir / f\"{target_id}_{msa_idx}.npz\"\n                msa_id_map[msa_id] = f\"{target_id}_{msa_idx}\"\n                if not processed.exists():\n                    # Parse A3M\n                    if msa_path.suffix == \".a3m\":\n                        msa: MSA = parse_a3m(\n                            msa_path,\n                            taxonomy=None,\n                            max_seqs=max_msa_seqs,\n                        )\n                    elif msa_path.suffix == \".csv\":\n                        msa: MSA = parse_csv(msa_path, max_seqs=max_msa_seqs)\n                    else:\n                        msg = f\"MSA file {msa_path} not supported, only a3m or csv.\"\n                        raise RuntimeError(msg)\n\n                    msa.dump(processed)\n\n            # Modify records to point to processed MSA\n            for c in target.record.chains:\n                if (c.msa_id != -1) and (c.msa_id in msa_id_map):\n                    c.msa_id = msa_id_map[c.msa_id]\n\n            # Keep record\n            records.append(target.record)\n\n            # Dump structure\n            struct_path = structure_dir / f\"{target.record.id}.npz\"\n            target.structure.dump(struct_path)\n\n        except Exception as e:\n            if len(data) > 1:\n                print(f\"Failed to process {path}. Skipping. Error: {e}.\")\n            else:\n                raise e\n\n    # Dump manifest\n    manifest = Manifest(records)\n    manifest.dump(out_dir / \"processed\" / \"manifest.json\")\n\ndef predict(\n    data: str,\n    out_dir: str,\n    cache: str = \"~/.boltz\",\n    checkpoint: Optional[str] = None,\n    devices: int = 1,\n    accelerator: str = \"gpu\",\n    recycling_steps: int = 3,\n    sampling_steps: int = 200,\n    diffusion_samples: int = 1,\n    step_scale: float = 1.638,\n    write_full_pae: bool = False,\n    write_full_pde: bool = False,\n    output_format: Literal[\"pdb\", \"mmcif\"] = \"mmcif\",\n    num_workers: int = 2,\n    override: bool = False,\n    seed: Optional[int] = None,\n    use_msa_server: bool = False,\n    msa_server_url: str = \"https://api.colabfold.com\",\n    msa_pairing_strategy: str = \"greedy\",\n) -> None:\n    \"\"\"Run predictions with Boltz-1.\"\"\"\n    # If cpu, write a friendly warning\n    if accelerator == \"cpu\":\n        msg = \"Running on CPU, this will be slow. Consider using a GPU.\"\n        click.echo(msg)\n\n    # Set no grad\n    torch.set_grad_enabled(False)\n\n    # Ignore matmul precision warning\n    torch.set_float32_matmul_precision(\"highest\")\n\n    # Set seed if desired\n    if seed is not None:\n        seed_everything(int(seed))\n\n    # Set cache path\n    cache = Path(cache).expanduser()\n    cache.mkdir(parents=True, exist_ok=True)\n\n    # Create output directories\n    data = Path(data).expanduser()\n    out_dir = Path(out_dir).expanduser()\n    out_dir = out_dir / f\"boltz_results_{data.stem}\"\n    out_dir.mkdir(parents=True, exist_ok=True)\n\n    # Download necessary data and model\n    download(cache)\n\n    # Validate inputs\n    data = check_inputs(data, out_dir, override)\n    if not data:\n        click.echo(\"No predictions to run, exiting.\")\n        return\n\n    # Set up trainer\n    strategy = \"auto\"\n    if (isinstance(devices, int) and devices > 1) or (\n        isinstance(devices, list) and len(devices) > 1\n    ):\n        strategy = DDPStrategy()\n        if len(data) < devices:\n            msg = (\n                \"Number of requested devices is greater \"\n                \"than the number of predictions.\"\n            )\n            raise ValueError(msg)\n\n    msg = f\"Running predictions for {len(data)} structure\"\n    msg += \"s\" if len(data) > 1 else \"\"\n    click.echo(msg)\n\n    # Process inputs\n    ccd_path = cache / \"ccd.pkl\"\n    process_inputs(\n        data=data,\n        out_dir=out_dir,\n        ccd_path=ccd_path,\n        use_msa_server=use_msa_server,\n        msa_server_url=msa_server_url,\n        msa_pairing_strategy=msa_pairing_strategy,\n    )\n\n    # Load processed data\n    processed_dir = out_dir / \"processed\"\n    processed = BoltzProcessedInput(\n        manifest=Manifest.load(processed_dir / \"manifest.json\"),\n        targets_dir=processed_dir / \"structures\",\n        msa_dir=processed_dir / \"msa\",\n    )\n\n    # Create data module\n    data_module = BoltzInferenceDataModule(\n        manifest=processed.manifest,\n        target_dir=processed.targets_dir,\n        msa_dir=processed.msa_dir,\n        num_workers=num_workers,\n    )\n\n    # Load model\n    if checkpoint is None:\n        checkpoint = cache / \"boltz1_conf.ckpt\"\n\n    predict_args = {\n        \"recycling_steps\": recycling_steps,\n        \"sampling_steps\": sampling_steps,\n        \"diffusion_samples\": diffusion_samples,\n        \"write_confidence_summary\": True,\n        \"write_full_pae\": write_full_pae,\n        \"write_full_pde\": write_full_pde,\n    }\n    diffusion_params = BoltzDiffusionParams()\n    diffusion_params.step_scale = step_scale\n    model_module: Boltz1 = Boltz1.load_from_checkpoint(\n        checkpoint,\n        strict=True,\n        predict_args=predict_args,\n        map_location=\"cpu\",\n        diffusion_process_args=asdict(diffusion_params),\n        ema=False,\n    )\n    model_module.eval()\n\n    # Create prediction writer\n    pred_writer = BoltzWriter(\n        data_dir=processed.targets_dir,\n        output_dir=out_dir / \"predictions\",\n        output_format=output_format,\n    )\n\n    trainer = Trainer(\n        default_root_dir=out_dir,\n        strategy=strategy,\n        callbacks=[pred_writer],\n        accelerator=accelerator,\n        devices=devices,\n        precision=32,\n    )\n\n    # Compute predictions\n    trainer.predict(\n        model_module,\n        datamodule=data_module,\n        return_predictions=False,\n    )\n\n\nif __name__ == \"__main__\":\n    predict(data=\"./inputs_prediction\",\n            out_dir=\"./outputs_prediction\",\n            cache=\"/kaggle/input/rna-prediction-boltz/\",\n            diffusion_samples=5,\n            seed=42,\n            override=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:39.823193Z","iopub.execute_input":"2025-03-07T15:44:39.823553Z","iopub.status.idle":"2025-03-07T15:44:39.831971Z","shell.execute_reply.started":"2025-03-07T15:44:39.823525Z","shell.execute_reply":"2025-03-07T15:44:39.830902Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prepare inputs","metadata":{}},{"cell_type":"code","source":"sub_file = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/test_sequences.csv')\n\nsub_file.head()\n\nnames = sub_file['target_id'].tolist()\nsequences = sub_file['sequence'].tolist()\n\n# Inference\nidx = 0 \nfor tmp_id, tmp_sequence in zip(names, sequences):\n    with open(f'/kaggle/working/inputs_prediction/{tmp_id}.yaml', 'w') as f:\n        f.write(\"constraints: []\\n\")\n        f.write(\"sequences:\\n\")\n        f.write(\"- rna:\\n\")\n        f.write(\"    id:\\n\")\n        f.write(\"    - A1\\n\")\n        f.write(f\"    sequence: {tmp_sequence}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:40.247667Z","iopub.execute_input":"2025-03-07T15:44:40.247952Z","iopub.status.idle":"2025-03-07T15:44:40.256754Z","shell.execute_reply.started":"2025-03-07T15:44:40.247930Z","shell.execute_reply":"2025-03-07T15:44:40.256039Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%ls inputs_prediction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:40.456754Z","iopub.execute_input":"2025-03-07T15:44:40.457060Z","iopub.status.idle":"2025-03-07T15:44:40.583533Z","shell.execute_reply.started":"2025-03-07T15:44:40.457035Z","shell.execute_reply":"2025-03-07T15:44:40.582500Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%ls outputs_prediction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:40.797337Z","iopub.execute_input":"2025-03-07T15:44:40.797662Z","iopub.status.idle":"2025-03-07T15:44:40.924304Z","shell.execute_reply.started":"2025-03-07T15:44:40.797638Z","shell.execute_reply":"2025-03-07T15:44:40.923180Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Exec inference","metadata":{}},{"cell_type":"code","source":"import torch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:41.457047Z","iopub.execute_input":"2025-03-07T15:44:41.457441Z","iopub.status.idle":"2025-03-07T15:44:41.461640Z","shell.execute_reply.started":"2025-03-07T15:44:41.457408Z","shell.execute_reply":"2025-03-07T15:44:41.460709Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.cuda.empty_cache()\nimport gc\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:43.443895Z","iopub.execute_input":"2025-03-07T15:44:43.444189Z","iopub.status.idle":"2025-03-07T15:44:43.567702Z","shell.execute_reply.started":"2025-03-07T15:44:43.444167Z","shell.execute_reply":"2025-03-07T15:44:43.566799Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import subprocess\nimport logging\n\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)\n\nresult = subprocess.run(['python', 'inference.py'], capture_output=True, text=True)\nlogger.info(f\"Command output: {result.stdout}\")\nlogger.error(f\"Command error: {result.stderr}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:44:43.622435Z","iopub.execute_input":"2025-03-07T15:44:43.622717Z","iopub.status.idle":"2025-03-07T15:46:10.129792Z","shell.execute_reply.started":"2025-03-07T15:44:43.622688Z","shell.execute_reply":"2025-03-07T15:46:10.129003Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Read RNA files","metadata":{}},{"cell_type":"code","source":"result","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:46:10.130829Z","iopub.execute_input":"2025-03-07T15:46:10.131113Z","iopub.status.idle":"2025-03-07T15:46:10.136079Z","shell.execute_reply.started":"2025-03-07T15:46:10.131088Z","shell.execute_reply":"2025-03-07T15:46:10.135335Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Gather results","metadata":{"execution":{"iopub.status.busy":"2025-03-04T21:51:26.943982Z","iopub.execute_input":"2025-03-04T21:51:26.94429Z","iopub.status.idle":"2025-03-04T21:51:26.949385Z","shell.execute_reply.started":"2025-03-04T21:51:26.944265Z","shell.execute_reply":"2025-03-04T21:51:26.948635Z"}}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio.PDB.MMCIF2Dict import MMCIF2Dict\n\ndef get_coords(tmp_id, idx):\n    cif_file = f\"outputs_prediction/boltz_results_inputs_prediction/predictions/{tmp_id}/{tmp_id}_model_{idx}.cif\"\n\n    mmcif_dict = MMCIF2Dict(cif_file)\n    \n    entity_poly_seq = mmcif_dict.get(\"_entity_poly_seq.mon_id\", [])\n    sequence = \"\".join(entity_poly_seq)\n    print(\"RNA sequence:\", sequence)\n    \n    x_coords = mmcif_dict[\"_atom_site.Cartn_x\"]\n    y_coords = mmcif_dict[\"_atom_site.Cartn_y\"]\n    z_coords = mmcif_dict[\"_atom_site.Cartn_z\"]\n    atom_names = mmcif_dict[\"_atom_site.label_atom_id\"]\n    \n    c1_coords = []\n    for i, atom in enumerate(atom_names):\n        if atom == \"C1'\":\n            c1_coords.append((float(x_coords[i]), float(y_coords[i]), float(z_coords[i])))\n    return c1_coords\n\nall_preds = os.listdir('outputs_prediction/boltz_results_inputs_prediction/predictions')\nsubmission = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/sample_submission.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T16:03:03.450269Z","iopub.execute_input":"2025-03-07T16:03:03.450615Z","iopub.status.idle":"2025-03-07T16:03:03.465813Z","shell.execute_reply.started":"2025-03-07T16:03:03.450589Z","shell.execute_reply":"2025-03-07T16:03:03.464703Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"idx = 0\nfor tmp_id in all_preds:\n    print('#'*20, f'inferences for {tmp_id}')\n    for idx in range(5):\n        c1_coords = get_coords(tmp_id, idx)\n        submission.loc[submission['ID'].apply(lambda x: tmp_id in x), [f'x_{idx+1}', f'y_{idx+1}', f'z_{idx+1}']] = c1_coords\n    print()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T16:04:42.502536Z","iopub.execute_input":"2025-03-07T16:04:42.502827Z","iopub.status.idle":"2025-03-07T16:04:42.881700Z","shell.execute_reply.started":"2025-03-07T16:04:42.502807Z","shell.execute_reply":"2025-03-07T16:04:42.880866Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.612270Z","iopub.execute_input":"2025-03-07T15:40:09.612483Z","iopub.status.idle":"2025-03-07T15:40:09.738568Z","shell.execute_reply.started":"2025-03-07T15:40:09.612464Z","shell.execute_reply":"2025-03-07T15:40:09.737492Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%rm -rf boltz","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.739730Z","iopub.execute_input":"2025-03-07T15:40:09.739995Z","iopub.status.idle":"2025-03-07T15:40:09.743674Z","shell.execute_reply.started":"2025-03-07T15:40:09.739960Z","shell.execute_reply":"2025-03-07T15:40:09.742874Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%rm -rf inputs_prediction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.739730Z","iopub.execute_input":"2025-03-07T15:40:09.739995Z","iopub.status.idle":"2025-03-07T15:40:09.743674Z","shell.execute_reply.started":"2025-03-07T15:40:09.739960Z","shell.execute_reply":"2025-03-07T15:40:09.742874Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%rm -rf outputs_prediction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.739730Z","iopub.execute_input":"2025-03-07T15:40:09.739995Z","iopub.status.idle":"2025-03-07T15:40:09.743674Z","shell.execute_reply.started":"2025-03-07T15:40:09.739960Z","shell.execute_reply":"2025-03-07T15:40:09.742874Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%rm -rf inference.py","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.739730Z","iopub.execute_input":"2025-03-07T15:40:09.739995Z","iopub.status.idle":"2025-03-07T15:40:09.743674Z","shell.execute_reply.started":"2025-03-07T15:40:09.739960Z","shell.execute_reply":"2025-03-07T15:40:09.742874Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.739730Z","iopub.execute_input":"2025-03-07T15:40:09.739995Z","iopub.status.idle":"2025-03-07T15:40:09.743674Z","shell.execute_reply.started":"2025-03-07T15:40:09.739960Z","shell.execute_reply":"2025-03-07T15:40:09.742874Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"submission.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.744516Z","iopub.execute_input":"2025-03-07T15:40:09.744742Z","iopub.status.idle":"2025-03-07T15:40:09.792630Z","shell.execute_reply.started":"2025-03-07T15:40:09.744721Z","shell.execute_reply":"2025-03-07T15:40:09.791995Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.793312Z","iopub.execute_input":"2025-03-07T15:40:09.793515Z","iopub.status.idle":"2025-03-07T15:40:09.919286Z","shell.execute_reply.started":"2025-03-07T15:40:09.793498Z","shell.execute_reply":"2025-03-07T15:40:09.918081Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission['target_id'] = submission['ID'].apply(lambda x: x.split('_')[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.920467Z","iopub.execute_input":"2025-03-07T15:40:09.920777Z","iopub.status.idle":"2025-03-07T15:40:09.927787Z","shell.execute_reply.started":"2025-03-07T15:40:09.920741Z","shell.execute_reply":"2025-03-07T15:40:09.927048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission.groupby('target_id')['x_1'].mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T15:40:09.928625Z","iopub.execute_input":"2025-03-07T15:40:09.928841Z","iopub.status.idle":"2025-03-07T15:40:09.953643Z","shell.execute_reply.started":"2025-03-07T15:40:09.928823Z","shell.execute_reply":"2025-03-07T15:40:09.952827Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}