diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml index 32fe61f..9dbfdb7 100644 --- a/.github/workflows/linting.yml +++ b/.github/workflows/linting.yml @@ -1,49 +1,57 @@ name: Linting -on: [push, pull_request] +on: + push: + branches: + - main + - master + pull_request: + branches: + - main + - master jobs: lint: runs-on: ubuntu-latest steps: - - name: Checkout - uses: actions/checkout@v4 - - - name: Set up Python 3.11 - uses: actions/setup-python@v4 - with: - python-version: '3.11' - - - name: Install uv - run: | - curl -LsSf https://astral.sh/uv/install.sh | sh - echo "$HOME/.cargo/bin" >> $GITHUB_PATH - - - name: Cache dependencies - uses: actions/cache@v4 - with: - path: | - ~/.cache/uv - ~/.uv - .venv - key: ${{ runner.os }}-uv-${{ hashFiles('pyproject.toml') }} - restore-keys: | - ${{ runner.os }}-uv- - - - name: Create and activate virtual environment - run: | - uv venv - echo "$PWD/.venv/bin" >> $GITHUB_PATH - - - name: Install dependencies - run: uv sync - - - name: Run ruff (linter) - run: ruff check - - - name: Run ruff (formatter) - run: ruff format --check - - - name: Run mypy - run: mypy . \ No newline at end of file + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v4 + with: + python-version: "3.11" + + - name: Install uv + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.cargo/bin" >> $GITHUB_PATH + + - name: Cache dependencies + uses: actions/cache@v4 + with: + path: | + ~/.cache/uv + ~/.uv + .venv + key: ${{ runner.os }}-uv-${{ hashFiles('pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-uv- + + - name: Create and activate virtual environment + run: | + uv venv + echo "$PWD/.venv/bin" >> $GITHUB_PATH + + - name: Install dependencies + run: uv sync + + - name: Run ruff (linter) + run: ruff check + + - name: Run ruff (formatter) + run: ruff format --check + + - name: Run mypy + run: mypy . diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 2196475..5dbaae1 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -1,42 +1,51 @@ name: Tests -on: [pull_request] +on: + push: + branches: + - main + - master + pull_request: + branches: + - main + - master + jobs: test: runs-on: ubuntu-latest steps: - - name: Checkout - uses: actions/checkout@v4 - - - name: Set up Python 3.13 - uses: actions/setup-python@v4 - with: - python-version: '3.13' - - - name: Install uv - run: | - curl -LsSf https://astral.sh/uv/install.sh | sh - echo "$HOME/.cargo/bin" >> $GITHUB_PATH - - - name: Cache dependencies - uses: actions/cache@v4 - with: - path: | - ~/.cache/uv - ~/.uv - .venv - key: ${{ runner.os }}-uv-${{ hashFiles('pyproject.toml') }} - restore-keys: | - ${{ runner.os }}-uv- - - - name: Create and activate virtual environment - run: | - uv venv - echo "$PWD/.venv/bin" >> $GITHUB_PATH - - - name: Install dependencies - run: uv sync - - - name: Run tests - run: pytest -v \ No newline at end of file + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install uv + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.cargo/bin" >> $GITHUB_PATH + + - name: Cache dependencies + uses: actions/cache@v4 + with: + path: | + ~/.cache/uv + ~/.uv + .venv + key: ${{ runner.os }}-uv-${{ hashFiles('pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-uv- + + - name: Create and activate virtual environment + run: | + uv venv + echo "$PWD/.venv/bin" >> $GITHUB_PATH + + - name: Install dependencies + run: uv sync + + - name: Run tests + run: pytest -v diff --git a/.gitignore b/.gitignore index 8855ca9..48f30c0 100644 --- a/.gitignore +++ b/.gitignore @@ -50,3 +50,5 @@ debug* slurm*.out submit*.sh qual.sh + +collectors/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3f0afee..060eb52 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ repos: entry: ruff check --fix language: python types: [python] - + - id: ruff-format name: ruff (formatter) entry: ruff format @@ -18,8 +18,8 @@ repos: - id: mypy name: mypy - entry: mypy + entry: uv run mypy language: system types: [python] exclude: ^(tests/) - args: ["--config-file=pyproject.toml"] \ No newline at end of file + args: ["--config-file=pyproject.toml"] diff --git a/docs/batched-generation.md b/docs/batched-generation.md new file mode 100644 index 0000000..c12e0bc --- /dev/null +++ b/docs/batched-generation.md @@ -0,0 +1,277 @@ +# Batched Route Generation + +The `BatchedBeamSearch` class provides efficient batched route generation for multiple target molecules simultaneously, with support for variable batch sizes and lengths. + +## Features + +- **Variable Batch Sizes**: Process any number of targets in a single batch +- **Variable Path Start Lengths**: Each target can have different starting material lengths +- **Variable Target Lengths**: Different maximum output lengths per target +- **Early Termination**: Each batch item can finish independently +- **GPU Efficient**: Optimized batching for maximum GPU utilization + +## Basic Usage + +### Single Target (Compatible with BeamSearchOptimized) + +```python +from pathlib import Path +from directmultistep import generate_routes + +target = "CNCc1ccccc1" +starting_material = "CN" +n_steps = 1 + +routes = generate_routes( + target=target, + n_steps=n_steps, + starting_material=starting_material, + beam_size=5, + model="flash", + config_path=Path("data/configs/dms_dictionary.yaml"), + ckpt_dir=Path("data/checkpoints"), +) + +for route in routes: + print(route) +``` + +### Multiple Targets (Batched) + +```python +from pathlib import Path +from directmultistep import generate_routes_batched + +targets = [ + "CNCc1ccccc1", + "CCOc1ccccc1", + "c1ccccc1", +] + +n_steps_list = [1, 2, 1] + +starting_materials = [ + "CN", + None, + None, +] + +routes = generate_routes_batched( + targets=targets, + n_steps_list=n_steps_list, + starting_materials=starting_materials, + beam_size=5, + model="flash", + config_path=Path("data/configs/dms_dictionary.yaml"), + ckpt_dir=Path("data/checkpoints"), +) + +for i, (target, routes_for_target) in enumerate(zip(targets, routes)): + print(f"Target {i+1}: {target}") + print(f"Routes: {len(routes_for_target)}") + for route in routes_for_target[:3]: + print(f" {route}") +``` + +## Advanced Usage + +### Using the Low-Level API + +For more control, you can use the lower-level APIs: + +```python +from pathlib import Path +import torch +from directmultistep import ( + load_published_model, + create_batched_beam_search, + prepare_batched_input_tensors, +) +from directmultistep.utils.dataset import RoutesProcessing + +# Load model +model = load_published_model("flash", Path("data/checkpoints")) +rds = RoutesProcessing(metadata_path=Path("data/configs/dms_dictionary.yaml")) + +# Create batched beam search +beam_search = create_batched_beam_search(model, beam_size=5, rds=rds) + +# Prepare batched inputs +targets = ["CNCc1ccccc1", "CCOc1ccccc1"] +n_steps_list = [1, 2] +starting_materials = ["CN", None] + +encoder_batch, steps_batch, path_starts, target_lengths = prepare_batched_input_tensors( + targets=targets, + n_steps_list=n_steps_list, + starting_materials=starting_materials, + rds=rds, + product_max_length=rds.product_max_length, + sm_max_length=rds.sm_max_length, +) + +# Run batched beam search +device = next(model.parameters()).device +results = beam_search.decode( + src_BC=encoder_batch.to(device), + steps_B1=steps_batch.to(device) if steps_batch is not None else None, + path_starts=[ps.to(device) for ps in path_starts], + target_lengths=target_lengths, + progress_bar=True, +) + +# Results is a list of lists: results[batch_idx][beam_idx] = (sequence, log_prob) +for batch_idx, beam_results in enumerate(results): + print(f"\nTarget {batch_idx}: {targets[batch_idx]}") + for beam_idx, (sequence, log_prob) in enumerate(beam_results): + print(f" Beam {beam_idx}: score={log_prob:.2f}, seq={sequence[:50]}...") +``` + +### Custom Batch Processing + +You can also directly use `BatchedBeamSearch` for custom processing: + +```python +from directmultistep.generation.tensor_gen import BatchedBeamSearch + +# Create custom beam search with specific parameters +beam_search = BatchedBeamSearch( + model=model, + beam_size=10, + start_idx=0, + pad_idx=52, + end_idx=22, + max_length=1074, + idx_to_token=rds.idx_to_token, + device=device, +) + +# Use with custom target lengths per batch item +results = beam_search.decode( + src_BC=encoder_batch, + steps_B1=steps_batch, + path_starts=path_starts, + target_lengths=[500, 1000, 1500], # Different max length per target + progress_bar=True, +) +``` + +## API Reference + +### High-Level Functions + +#### `generate_routes_batched` + +```python +def generate_routes_batched( + targets: Sequence[str], + n_steps_list: Sequence[int] | None, + starting_materials: Sequence[str | None], + beam_size: int, + model: ModelName | torch.nn.Module, + config_path: Path, + ckpt_dir: Path | None = None, + commercial_stock: set[str] | None = None, + use_fp16: bool = False, +) -> list[list[str]]: +``` + +Generate synthesis routes for multiple targets using batched beam search. + +**Arguments:** +- `targets`: List of SMILES strings of target molecules +- `n_steps_list`: List of number of synthesis steps for each target or None (for explorer) +- `starting_materials`: List of starting materials for each target (can contain None) +- `beam_size`: Beam size for the beam search +- `model`: Either a model name or a torch.nn.Module +- `config_path`: Path to the model configuration file +- `ckpt_dir`: Directory containing model checkpoints (required if model is a string) +- `commercial_stock`: Set of commercially available starting materials (SMILES) +- `use_fp16`: Whether to use half precision (FP16) + +**Returns:** +- List of lists, where each inner list contains valid routes for the corresponding target + +### Utility Functions + +#### `prepare_batched_input_tensors` + +```python +def prepare_batched_input_tensors( + targets: Sequence[str], + n_steps_list: Sequence[int] | None, + starting_materials: Sequence[str | None], + rds: RoutesProcessing, + product_max_length: int, + sm_max_length: int, + use_fp16: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor], list[int]]: +``` + +Prepare batched input tensors for the model. + +**Returns:** +- `encoder_batch`: Batched input tensor for the encoder [B, C] +- `steps_batch`: Batched tensor of steps [B, 1], or None if all n_steps are None +- `path_starts`: List of initial path tensors for decoder (variable lengths) +- `target_lengths`: List of target max lengths per batch item + +#### `create_batched_beam_search` + +```python +def create_batched_beam_search( + model: torch.nn.Module, + beam_size: int, + rds: RoutesProcessing +) -> BatchedBeamSearch: +``` + +Create a batched beam search object that supports variable batch sizes and lengths. + +### BatchedBeamSearch Class + +```python +class BatchedBeamSearch: + def __init__( + self, + model: nn.Module, + beam_size: int, + start_idx: int, + pad_idx: int, + end_idx: int, + max_length: int, + idx_to_token: dict[int, str], + device: torch.device, + ): + ... + + def decode( + self, + src_BC: Tensor, + steps_B1: Tensor | None, + path_starts: list[Tensor | None] | None = None, + target_lengths: list[int] | None = None, + progress_bar: bool = True, + token_processor: Callable[[list[str]], str] | None = None, + ) -> list[list[tuple[str, float]]]: + ... +``` + +## Performance Considerations + +1. **Batch Size**: Larger batches improve GPU utilization but increase memory usage +2. **Variable Lengths**: The implementation handles variable lengths efficiently by grouping active beams +3. **Early Termination**: Batch items that finish early are removed from computation +4. **Memory Usage**: Peak memory scales with `batch_size * beam_size * max_sequence_length` + +## Comparison with BeamSearchOptimized + +| Feature | BeamSearchOptimized | BatchedBeamSearch | +|---------|-------------------|------------------| +| Batch Size | Only 1 | Any positive integer | +| Variable Path Starts | No | Yes | +| Variable Target Lengths | No | Yes | +| Early Termination | All beams together | Per batch item | +| API Compatibility | Single target | Multiple targets | + +For single-target generation, both implementations produce identical results. For multiple targets, use `BatchedBeamSearch` for better efficiency. diff --git a/examples/batched_generation_example.py b/examples/batched_generation_example.py new file mode 100644 index 0000000..ad1641f --- /dev/null +++ b/examples/batched_generation_example.py @@ -0,0 +1,46 @@ +""" +Example script demonstrating batched route generation. + +This shows how to use the BatchedBeamSearch class to generate routes +for multiple targets simultaneously with different starting materials +and step counts. +""" + +from pathlib import Path + +from directmultistep.generate import generate_routes_batched + +config_path = Path("data/configs/dms_dictionary.yaml") +ckpt_dir = Path("data/checkpoints") + +targets = [ + "CNCc1ccccc1", + "CCOc1ccccc1", + "c1ccccc1", +] + +n_steps_list = [1, 2, 1] + +starting_materials = [ + "CN", + None, + None, +] + +routes = generate_routes_batched( + targets=targets, + n_steps_list=n_steps_list, + starting_materials=starting_materials, + beam_size=5, + model="flash", + config_path=config_path, + ckpt_dir=ckpt_dir, +) + +for i, (target, routes_for_target) in enumerate(zip(targets, routes, strict=False)): + print(f"\nTarget {i + 1}: {target}") + print(f"Starting material: {starting_materials[i]}") + print(f"Number of steps: {n_steps_list[i]}") + print(f"Generated {len(routes_for_target)} valid routes:") + for j, route in enumerate(routes_for_target[:3], 1): + print(f" Route {j}: {route[:100]}..." if len(route) > 100 else f" Route {j}: {route}") diff --git a/pyproject.toml b/pyproject.toml index 202d33d..4091fed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ dependencies = [ "rdkit>=2023.9.3", "torch>=2.3.0", "tqdm>=4.67.1", - + ] authors = [ { name = "Anton Morgunov", email = "anton@ischemist.com" }, @@ -80,6 +80,7 @@ lint.select = [ "SIM", # flake8-simplify "I", # isort ] -lint.ignore = ["E501"] - +lint.ignore = ["E501"] +[tool.pytest.ini_options] +addopts = "-m 'not ckptreq'" diff --git a/scripts/dev/run-beam.py b/scripts/dev/run-beam.py new file mode 100644 index 0000000..e2d7a81 --- /dev/null +++ b/scripts/dev/run-beam.py @@ -0,0 +1,97 @@ +from pathlib import Path + +from directmultistep import generate_routes, generate_routes_batched + + +def run_beam1() -> None: + target = "CNCc1ccccc1" + starting_material = "CN" + n_steps = 1 + + routes = generate_routes( + target=target, + n_steps=n_steps, + starting_material=starting_material, + beam_size=5, + model="flash", + config_path=Path("data/configs/dms_dictionary.yaml"), + ckpt_dir=Path("data/checkpoints"), + ) + + with open("b-1-routes.txt", "w") as f: + for route in routes[:3]: + f.write(route + "\n") + + +def run_beam2() -> None: + targets = ["CNCc1ccccc1"] * 2 + starting_materials = ["CN"] * 2 + n_steps_list = [1] * 2 + + routes = generate_routes_batched( + targets=targets, + n_steps_list=n_steps_list, + starting_materials=starting_materials, + beam_size=5, + model="flash", + config_path=Path("data/configs/dms_dictionary.yaml"), + ckpt_dir=Path("data/checkpoints"), + ) + + with open("b-2-routes.txt", "w") as f: + for i, (target, routes_for_target) in enumerate(zip(targets, routes, strict=False)): + f.write(f"Target {i + 1}: {target}\n") + f.write(f"Routes: {len(routes_for_target)}\n") + for route in routes_for_target[:3]: + f.write(route + "\n") + + +def run_beam_hard() -> None: + targets_list = [ + "CC(=O)OC1=CC=CC=C1C(=O)O", + "CNCc1cc(-c2ccccc2F)n(S(=O)(=O)c2cccnc2)c1", + "O=C(c1ccc(NS(=O)(=O)c2cccc3cccnc23)cc1)N1CCN(CC2CC2)CC1", + "COc1ccc(-n2nccn2)c(C(=O)N2CCC[C@@]2(C)c2nc3c(C)c(Cl)ccc3[nH]2)c1", + ] + sms_list = [None, "O=S(=O)(Cl)c1cccnc1", "CCOC(=O)c1ccc(N)cc1", "C[C@@]1(C(=O)O)CCCN1"] + n_steps_list = [1, 2, 5, 4] + + from tqdm import tqdm + + beams = [5, 20, 50] + + for beam in beams: + print(f"Beam size: {beam}") + for fp_16 in [False, True]: + print(f"Use fp16: {fp_16}") + for target, sm, n_steps in tqdm( + zip(targets_list, sms_list, n_steps_list, strict=False), total=len(targets_list) + ): + generate_routes( + target=target, + n_steps=n_steps, + starting_material=sm, + beam_size=beam, + model="flash", + config_path=Path("data/configs/dms_dictionary.yaml"), + ckpt_dir=Path("data/checkpoints"), + show_progress=False, + use_fp16=fp_16, + ) + + generate_routes_batched( + targets=targets_list, + n_steps_list=n_steps_list, + starting_materials=sms_list, + beam_size=beam, + model="flash", + config_path=Path("data/configs/dms_dictionary.yaml"), + ckpt_dir=Path("data/checkpoints"), + use_fp16=fp_16, + ) + + +if __name__ == "__main__": + run_beam1() + run_beam2() + run_beam_hard() diff --git a/src/directmultistep/__init__.py b/src/directmultistep/__init__.py index 1a6b3f4..08e7e72 100644 --- a/src/directmultistep/__init__.py +++ b/src/directmultistep/__init__.py @@ -1,5 +1,27 @@ """DirectMultiStep - Direct Route Generation for Multi-Step Retrosynthesis.""" +from directmultistep.generate import ( + create_batched_beam_search, + create_beam_search, + generate_routes, + generate_routes_batched, + load_published_model, + prepare_batched_input_tensors, + prepare_input_tensors, +) +from directmultistep.generation.tensor_gen import BatchedBeamSearch, BeamSearchOptimized from directmultistep.utils.logging_config import setup_logging setup_logging() + +__all__ = [ + "BeamSearchOptimized", + "BatchedBeamSearch", + "create_batched_beam_search", + "create_beam_search", + "generate_routes", + "generate_routes_batched", + "load_published_model", + "prepare_batched_input_tensors", + "prepare_input_tensors", +] diff --git a/src/directmultistep/generate.py b/src/directmultistep/generate.py index 71dd0e2..cefd8c7 100644 --- a/src/directmultistep/generate.py +++ b/src/directmultistep/generate.py @@ -1,9 +1,11 @@ +from collections.abc import Sequence from pathlib import Path from typing import Literal, cast import torch import torch.nn as nn +from directmultistep.generation.tensor_gen import BatchedBeamSearch from directmultistep.generation.tensor_gen import BeamSearchOptimized as BeamSearch from directmultistep.model import ModelFactory from directmultistep.utils.dataset import RoutesProcessing @@ -73,6 +75,23 @@ def create_beam_search(model: torch.nn.Module, beam_size: int, rds: RoutesProces return beam +def create_batched_beam_search(model: torch.nn.Module, beam_size: int, rds: RoutesProcessing) -> BatchedBeamSearch: + """Create a batched beam search object that supports variable batch sizes and lengths.""" + device = next(model.parameters()).device + + beam = BatchedBeamSearch( + model=model, + beam_size=beam_size, + start_idx=0, + pad_idx=52, + end_idx=22, + max_length=1074, + idx_to_token=rds.idx_to_token, + device=device, + ) + return beam + + def prepare_input_tensors( target: str, n_steps: int | None, @@ -98,11 +117,12 @@ def prepare_input_tensors( - steps_tens: Tensor of the number of steps, or None if not provided. - path_tens: Initial path tensor for the decoder. """ - prod_tens = rds.smile_to_tokens(target, product_max_length) if starting_material: + prod_tens = rds.smile_to_tokens(target, product_max_length) sm_tens = rds.smile_to_tokens(starting_material, sm_max_length) encoder_inp = torch.cat([prod_tens, sm_tens], dim=0).unsqueeze(0) else: + prod_tens = rds.smile_to_tokens(target, product_max_length + sm_max_length) encoder_inp = torch.cat([prod_tens], dim=0).unsqueeze(0) steps_tens = torch.tensor([n_steps]).unsqueeze(0) if n_steps is not None else None @@ -118,6 +138,67 @@ def prepare_input_tensors( return encoder_inp, steps_tens, path_tens +def prepare_batched_input_tensors( + targets: Sequence[str], + n_steps_list: Sequence[int] | None, + starting_materials: Sequence[str | None], + rds: RoutesProcessing, + product_max_length: int, + sm_max_length: int, + use_fp16: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor], list[int]]: + """Prepare batched input tensors for the model. + + Args: + targets: List of SMILES strings of target molecules + n_steps_list: List of number of synthesis steps for each target, or None if not using steps + starting_materials: List of SMILES strings of starting materials (can contain None) + rds: RoutesProcessing object for tokenization + product_max_length: Maximum length of the product SMILES sequence + sm_max_length: Maximum length of the starting material SMILES sequence + use_fp16: Whether to use half precision (FP16) for tensors + + Returns: + A tuple containing: + - encoder_batch: Batched input tensor for the encoder [B, C] + - steps_batch: Batched tensor of steps [B, 1], or None if n_steps_list is None + - path_starts: List of initial path tensors for decoder (variable lengths) + - target_lengths: List of target max lengths per batch item + """ + if n_steps_list is not None and len(targets) != len(n_steps_list): + raise ValueError(f"Length mismatch: targets={len(targets)}, n_steps_list={len(n_steps_list)}") + if len(targets) != len(starting_materials): + raise ValueError(f"Length mismatch: targets={len(targets)}, starting_materials={len(starting_materials)}") + + encoder_inputs = [] + steps_tensors: list[torch.Tensor] = [] + path_starts = [] + target_lengths = [] + + for i, (target, sm) in enumerate(zip(targets, starting_materials, strict=False)): + n_steps = n_steps_list[i] if n_steps_list is not None else None + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + target=target, + n_steps=n_steps, + starting_material=sm, + rds=rds, + product_max_length=product_max_length, + sm_max_length=sm_max_length, + use_fp16=use_fp16, + ) + + encoder_inputs.append(encoder_inp.squeeze(0)) + if steps_tens is not None: + steps_tensors.append(steps_tens.squeeze(0)) + path_starts.append(path_tens.squeeze(0)) + target_lengths.append(1074) + + encoder_batch = torch.stack(encoder_inputs) + steps_batch = torch.stack(steps_tensors) if n_steps_list is not None else None + + return encoder_batch, steps_batch, path_starts, target_lengths + + def generate_routes( target: str, n_steps: int | None, @@ -128,6 +209,7 @@ def generate_routes( ckpt_dir: Path | None = None, commercial_stock: set[str] | None = None, use_fp16: bool = False, + show_progress: bool = True, ) -> list[str]: """Generate synthesis routes using the model. @@ -164,6 +246,7 @@ def generate_routes( src_BC=encoder_inp.to(device), steps_B1=steps_tens.to(device) if steps_tens is not None else None, path_start_BL=path_tens.to(device), + progress_bar=show_progress, ) for beam_result_S2 in beam_result_BS2: all_beam_results_NS2.append(beam_result_S2) @@ -177,3 +260,74 @@ def generate_routes( commercial_stock=commercial_stock, ) return [beam_result[0] for beam_result in correct_paths_NS2n[0]] + + +def generate_routes_batched( + targets: Sequence[str], + n_steps_list: Sequence[int] | None, + starting_materials: Sequence[str | None], + beam_size: int, + model: ModelName | torch.nn.Module, + config_path: Path, + ckpt_dir: Path | None = None, + commercial_stock: set[str] | None = None, + use_fp16: bool = False, +) -> list[list[str]]: + """Generate synthesis routes for multiple targets using batched beam search. + + Args: + targets: List of SMILES strings of target molecules + n_steps_list: List of number of synthesis steps for each target, or None if not using steps + starting_materials: List of starting materials for each target (can contain None) + beam_size: Beam size for the beam search + model: Either a model name or a torch.nn.Module + config_path: Path to the model configuration file + ckpt_dir: Directory containing model checkpoints (required if model is a string) + commercial_stock: Set of commercially available starting materials (SMILES) + use_fp16: Whether to use half precision (FP16) for model weights and computations + + Returns: + List of lists, where each inner list contains valid routes for the corresponding target + """ + if isinstance(model, str): + if ckpt_dir is None: + raise ValueError("ckpt_dir must be provided when model is specified by name") + for i, (sm, _target) in enumerate(zip(starting_materials, targets, strict=False)): + n_steps = n_steps_list[i] if n_steps_list is not None else None + validate_model_constraints(model, n_steps, sm) + model = load_published_model(model, ckpt_dir, use_fp16) + + rds = RoutesProcessing(metadata_path=config_path) + beam_obj = create_batched_beam_search(model, beam_size, rds) + + encoder_batch, steps_batch, path_starts, target_lengths = prepare_batched_input_tensors( + targets=targets, + n_steps_list=n_steps_list, + starting_materials=starting_materials, + rds=rds, + product_max_length=rds.product_max_length, + sm_max_length=rds.sm_max_length, + use_fp16=use_fp16, + ) + + device = next(model.parameters()).device + beam_results = beam_obj.decode( + src_BC=encoder_batch.to(device), + steps_B1=steps_batch.to(device) if steps_batch is not None else None, + path_starts=[ps.to(device) for ps in path_starts], + target_lengths=target_lengths, + progress_bar=True, + ) + + all_results = [] + for idx, (target, sm) in enumerate(zip(targets, starting_materials, strict=False)): + valid_paths = find_valid_paths([beam_results[idx]]) + correct_paths = process_path_single( + paths_NS2n=valid_paths, + true_products=[target], + true_reacs=[sm] if sm else None, + commercial_stock=commercial_stock, + ) + all_results.append([beam_result[0] for beam_result in correct_paths[0]]) + + return all_results diff --git a/src/directmultistep/generation/tensor_gen.py b/src/directmultistep/generation/tensor_gen.py index d13b26b..ec62f98 100644 --- a/src/directmultistep/generation/tensor_gen.py +++ b/src/directmultistep/generation/tensor_gen.py @@ -22,11 +22,272 @@ from directmultistep.utils.logging_config import logger -# Define types Tensor = torch.Tensor BeamSearchOutput = list[list[tuple[str, float]]] +class BatchedBeamSearch: + def __init__( + self, + model: nn.Module, + beam_size: int, + start_idx: int, + pad_idx: int, + end_idx: int, + max_length: int, + idx_to_token: dict[int, str], + device: torch.device, + ): + self.model = model + self.beam_size = beam_size + self.start_idx = start_idx + self.pad_idx = pad_idx + self.end_idx = end_idx + self.device = device + self.max_length = max_length + self.idx_to_token = idx_to_token + + # Pre-allocate reusable tensors + self._init_reusable_tensors() + + def __repr__(self) -> str: + return f"BatchedBeamSearch(beam_size={self.beam_size}, max_length={self.max_length})" + + def _init_reusable_tensors(self) -> None: + """Pre-allocate tensors that can be reused across decode calls""" + S = self.beam_size + # Pre-compute beam indices for gather operations + self.beam_indices = torch.arange(S, device=self.device) + + def decode( + self, + src_BC: Tensor, + steps_B1: Tensor | None, + path_starts: list[Tensor | None] | None = None, + target_lengths: list[int] | None = None, + progress_bar: bool = True, + token_processor: Callable[[list[str]], str] | None = None, + ) -> BeamSearchOutput: + """ + Optimized fully vectorized beam search. + """ + B, C = src_BC.shape + S = self.beam_size + L = self.max_length + # V = len(self.idx_to_token) + + # Detect the dtype from input (or could check model.parameters()) + dtype = src_BC.dtype + + # Encode source sequences once + src_mask_B11C = (src_BC != self.pad_idx).unsqueeze(1).unsqueeze(2) + enc_src_BCD = self.model.encoder(src_BC.long(), src_mask_B11C, steps_B1) + + # Optimization 1: Use view instead of reshape for encoder expansion + enc_src_BSCD = enc_src_BCD.unsqueeze(1).expand(B, S, -1, -1).contiguous().view(B * S, -1, enc_src_BCD.size(-1)) + src_mask_BS11C = src_mask_B11C.unsqueeze(1).expand(B, S, -1, -1, -1).contiguous().view(B * S, 1, 1, C) + + # Initialize beam tensors + sequences_BSL = torch.full((B, S, L), self.pad_idx, dtype=torch.long, device=self.device) + scores_BS = torch.full((B, S), float("-inf"), dtype=dtype, device=self.device) + scores_BS[:, 0] = 0.0 + active_BS = torch.ones((B, S), dtype=torch.bool, device=self.device) + + # Handle path starts + first_steps_B = torch.ones(B, dtype=torch.long, device=self.device) + if path_starts: + for b, path in enumerate(path_starts): + if path is not None: + path_len = path.size(0) + sequences_BSL[b, :, :path_len] = path.unsqueeze(0).expand(S, -1) + first_steps_B[b] = path_len + else: + sequences_BSL[b, :, 0] = self.start_idx + else: + sequences_BSL[:, :, 0] = self.start_idx + + first_step = int(first_steps_B.min().item()) + max_steps = L - 1 if target_lengths is None else max(target_lengths) + + # Pre-allocate candidate selection tensors + candidate_buffer = torch.zeros((B, S * S), device=self.device, dtype=dtype) + beam_idx_buffer = torch.zeros((B, S * S), dtype=torch.long, device=self.device) + token_idx_buffer = torch.zeros((B, S * S), dtype=torch.long, device=self.device) + + pbar: Iterable[int] = ( + tqdm(range(first_step, max_steps), desc="Beam search", dynamic_ncols=True) + if progress_bar + else range(first_step, max_steps) + ) + # OPTIMIZATION: Pre-compute beam patterns once (outside loop ideally) + beam_idx_pattern = torch.arange(S, device=self.device).unsqueeze(1).expand(S, S).reshape(-1) + + for step in pbar: + # Check which batches have started decoding + step_active_B = step >= first_steps_B + + # Check for end tokens + has_end_token_BS = (sequences_BSL[:, :, :step] == self.end_idx).any(dim=2) + active_BS &= ~has_end_token_BS + + # Check if all beams for each batch are inactive + batch_finished_B = ~active_BS.any(dim=1) | ~step_active_B + + if batch_finished_B.all(): + break + + if progress_bar and batch_finished_B.any(): + finished_count = batch_finished_B.sum().item() + pbar.set_postfix({"Finished batches": f"{finished_count}/{B}"}) # type: ignore + + # Create mask for active beams + active_mask_BS = active_BS & step_active_B.unsqueeze(1) + + if not active_mask_BS.any(): + continue + + # KEEP ORIGINAL: Only process active sequences + active_mask_flat = active_mask_BS.view(B * S) + + # Get active sequence indices + active_indices = active_mask_flat.nonzero(as_tuple=True)[0] + + # Process only active sequences + sequences_flat = sequences_BSL.view(B * S, L) + active_sequences = sequences_flat[active_indices, :step] + active_enc_src = enc_src_BSCD[active_indices] + active_src_mask = src_mask_BS11C[active_indices] + + # Forward pass only for active beams + with torch.no_grad(): + active_output = self.model.decoder( + trg_BL=active_sequences, + enc_src_BCD=active_enc_src, + src_mask_B11C=active_src_mask, + trg_mask_B1LL=None, + ) + + # Get log probabilities + active_log_probs = torch.log_softmax(active_output[:, -1, :], dim=-1) + + # Scatter back to full tensor + log_probs_BSV = torch.full( + (B * S, active_log_probs.size(-1)), float("-inf"), dtype=dtype, device=self.device + ) + log_probs_BSV[active_indices] = active_log_probs + log_probs_BSV = log_probs_BSV.view(B, S, -1) + + # Create expansion mask for first step logic + is_first_step_B = first_steps_B == step + expand_mask_BS = torch.ones((B, S), dtype=torch.bool, device=self.device) + expand_mask_BS[is_first_step_B, 1:] = False + + # Combined mask for beams that should generate candidates + generate_mask_BS = active_mask_BS & expand_mask_BS + + # Get top-k tokens for each beam + top_k_scores_BSS, top_k_tokens_BSS = torch.topk(log_probs_BSV, S, dim=-1) + + # Calculate candidate scores + candidate_scores_BSS = scores_BS.unsqueeze(-1) + top_k_scores_BSS + + # For first step beams, only use first beam's candidates + if is_first_step_B.any(): + first_step_mask = is_first_step_B.view(-1, 1, 1) + first_beam_scores = candidate_scores_BSS[:, 0:1, :].expand(B, S, S) + candidate_scores_BSS = torch.where(first_step_mask, first_beam_scores, candidate_scores_BSS) + + # Mask out invalid candidates + generate_mask_BSS = generate_mask_BS.unsqueeze(-1) + candidate_scores_BSS = candidate_scores_BSS.masked_fill(~generate_mask_BSS, float("-inf")) + + # Flatten candidates + candidate_buffer = candidate_scores_BSS.view(B, S * S) + beam_idx_buffer = beam_idx_pattern.unsqueeze(0).expand(B, -1) + token_idx_buffer = top_k_tokens_BSS.view(B, S * S) + + # Handle inactive beams + inactive_mask_BS = ~active_BS & (scores_BS > float("-inf")) + + # OPTIMIZATION: Use pre-allocated tensors and efficient concatenation + n_candidates = S * S + inactive_scores = torch.where(inactive_mask_BS, scores_BS, scores_BS.new_full((), float("-inf"))) + all_scores = torch.cat([candidate_buffer, inactive_scores], dim=1) + + all_beam_indices = torch.cat( + [beam_idx_buffer, torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1)], dim=1 + ) + + all_token_indices = torch.cat( + [token_idx_buffer, torch.full((B, S), -1, dtype=torch.long, device=self.device)], dim=1 + ) + + # Compute sequence lengths for normalization + seq_lengths_BS = (sequences_BSL != self.pad_idx).sum(dim=-1).float() + + # OPTIMIZATION: More efficient length computation + all_seq_lengths = torch.cat( + [seq_lengths_BS.unsqueeze(-1).expand(B, S, S).reshape(B, n_candidates) + 1, seq_lengths_BS], dim=1 + ) + + # Normalize + normalized_scores = all_scores / (all_seq_lengths.sqrt() + 1e-6) + + # Select top S beams for each batch + _, top_indices = torch.topk(normalized_scores, S, dim=1) + + # Gather selected beams + selected_beam_indices = torch.gather(all_beam_indices, 1, top_indices) + selected_token_indices = torch.gather(all_token_indices, 1, top_indices) + selected_scores = torch.gather(all_scores, 1, top_indices) + + # Update sequences + gather_indices = selected_beam_indices.unsqueeze(-1).expand(B, S, L) + sequences_BSL = torch.gather(sequences_BSL, 1, gather_indices) + + # Add new tokens where applicable + mask_add_token = selected_token_indices != -1 + batch_coords, beam_coords = torch.where(mask_add_token) + if batch_coords.numel() > 0: + sequences_BSL[batch_coords, beam_coords, step] = selected_token_indices[batch_coords, beam_coords] + + # Update active status + has_end_in_selected = (selected_token_indices == self.end_idx) | ( + sequences_BSL[:, :, :step] == self.end_idx + ).any(dim=2) + was_inactive = selected_token_indices == -1 + + # Update all states + scores_BS = selected_scores + active_BS = ~has_end_in_selected & ~was_inactive + + # Extract final results (unchanged) + outputs_BS2_nt: list[list[tuple[str, float]]] = [] + + for b in range(B): + batch_results = [] + for s in range(S): + if scores_BS[b, s] == float("-inf"): + continue + + output_tokens = [] + for idx in sequences_BSL[b, s]: + if idx == self.pad_idx: + break + if idx == self.start_idx: + continue + if idx == self.end_idx: + break + output_tokens.append(self.idx_to_token[idx.item()]) + + output_str = token_processor(output_tokens) if token_processor else "".join(output_tokens) + batch_results.append((output_str, scores_BS[b, s].item())) + + outputs_BS2_nt.append(batch_results) + + return outputs_BS2_nt + + class BeamSearchOptimized: def __init__( self, @@ -71,10 +332,9 @@ def decode( S = self.beam_size L = self.max_length - # Prepare mask and encoder outputs src_mask_B11C = (src_BC != self.pad_idx).unsqueeze(1).unsqueeze(2) enc_src_BCD = self.model.encoder(src_BC.long(), src_mask_B11C, steps_B1) - beam_enc_WCD = enc_src_BCD.repeat_interleave(S, dim=0) # W = B * S + beam_enc_WCD = enc_src_BCD.repeat_interleave(S, dim=0) beam_src_WC = src_BC.repeat_interleave(S, dim=0) beam_src_mask_W11C = (beam_src_WC != self.pad_idx).unsqueeze(1).unsqueeze(2) @@ -93,25 +353,25 @@ def decode( logger.info( f"Generating routes with beam size {S}. The progress bar may end early if all beams find end token." ) - pbar: Iterable[int] = tqdm(range(first_step, L - 1)) if progress_bar else range(first_step, L - 1) + pbar: Iterable[int] = ( + tqdm(range(first_step, L - 1), dynamic_ncols=True) if progress_bar else range(first_step, L - 1) + ) for step in pbar: with torch.no_grad(): output_WLV = self.model.decoder( trg_BL=beam_idxs_WL[:, :step], enc_src_BCD=beam_enc_WCD, src_mask_B11C=beam_src_mask_W11C, - trg_mask_B1LL=None, # trg_mask_W1LL[:, :, :step, :step] + trg_mask_B1LL=None, ) W, _, V = output_WLV.shape - output_WV = output_WLV[:, -1, :] # Get the last token's logits + output_WV = output_WLV[:, -1, :] log_probs_WV = torch.log_softmax(output_WV, dim=-1) finished_sequences_W = torch.any(beam_idxs_WL == self.end_idx, dim=-1) active_mask_W = ~finished_sequences_W if finished_sequences_W.all(): break - # finished_mask_WV = finished_sequences_W.unsqueeze(-1).expand(-1, V) - # log_probs_WV = log_probs_WV.masked_fill(finished_mask_WV, float('-inf')) if step == first_step: log_probs_BSV = log_probs_WV.view(B, S, -1) @@ -131,10 +391,9 @@ def decode( _S = active_beams_BSL.size(1) active_beams_BSSL = active_beams_BSL.unsqueeze(2).repeat(1, 1, S, 1) active_beams_BSSL[..., step] = act_top_k_idxs_BSS - active_beams_BSsqL = active_beams_BSSL.view(B, -1, L) # my candidate_seqs_BSL_nt + active_beams_BSsqL = active_beams_BSSL.view(B, -1, L) cur_log_probs_WS = cur_log_probs_WV[active_mask_W].view(-1, V).gather(1, act_top_k_idxs_WS) - # cur_log_probs_BSsq = cur_log_probs_WS.view(B, -1) # my candidate_probs_BS_nt sequence_lengths_WL = (active_beams_WL.ne(self.pad_idx).sum(dim=1).float()).unsqueeze(1) normalized_act_log_probs_WS = cur_log_probs_WS / (sequence_lengths_WL.sqrt() + 1e-6) diff --git a/src/directmultistep/model/factory.py b/src/directmultistep/model/factory.py index c607c2b..64cecac 100644 --- a/src/directmultistep/model/factory.py +++ b/src/directmultistep/model/factory.py @@ -16,6 +16,7 @@ Seq2SeqConfig, TransformerConfig, ) +from directmultistep.utils.logging_config import logger class ModelFactory: @@ -167,7 +168,7 @@ def create_model(self) -> Seq2Seq: if self.compile_model: model = torch.compile(model) # type: ignore - print(f"The model has {self._count_parameters(model):,} trainable parameters") + logger.debug(f"The model has {self._count_parameters(model):,} trainable parameters") return model @classmethod diff --git a/tests/generation/conftest.py b/tests/generation/conftest.py new file mode 100644 index 0000000..fc2190e --- /dev/null +++ b/tests/generation/conftest.py @@ -0,0 +1,118 @@ +""" +Pytest configuration and fixtures for generation tests. +""" + +from pathlib import Path + +import numpy as np +import pytest +import torch + +# Set seeds for reproducible tests +torch.manual_seed(42) +np.random.seed(42) + +# Test cases for beam search testing +TEST_CASES = [ + { + "name": "target1", + "target": "CNCc1cc(-c2ccccc2F)n(S(=O)(=O)c2cccnc2)c1", + "starting_material": "CN", + "n_steps": 2, + "description": "Simple target with primary amine starting material", + }, + { + "name": "target2", + "target": "O=C(c1ccc(NS(=O)(=O)c2cccc3cccnc23)cc1)N1CCN(CC2CC2)CC1", + "starting_material": "CCOC(=O)c1ccc(N)cc1", + "n_steps": 5, + "description": "Complex target with multiple functional groups", + }, + { + "name": "target3", + "target": "CC(C)(C)[C@@H](CS(C)(=O)=O)Nc1nc(-c2c[nH]c3ncc(F)cc23)ncc1F", + "starting_material": "CC(C)(C)[C@H](N)CO", + "n_steps": 3, + "description": "Complex pharmaceutical-like molecule", + }, +] + +SIMPLE_SMILES_CASES = [ + {"name": "simple_alkane", "smiles": "CCCC", "description": "Simple butane molecule"}, + {"name": "simple_arene", "smiles": "c1ccccc1", "description": "Benzene ring"}, + {"name": "simple_functional", "smiles": "CC(=O)O", "description": "Acetic acid"}, +] + + +@pytest.fixture(scope="session") +def test_cases(): + """Return the standard test cases for generation testing.""" + return TEST_CASES + + +@pytest.fixture(scope="session") +def simple_smiles_cases(): + """Return simple SMILES cases for basic testing.""" + return SIMPLE_SMILES_CASES + + +@pytest.fixture +def model_files_available(): + """Check if model files are available for testing.""" + config_path = Path("data/configs/dms_dictionary.yaml") + ckpt_dir = Path("data/checkpoints") + + return config_path.exists() and ckpt_dir.exists() and any(ckpt_dir.iterdir()) + + +@pytest.fixture +def test_data_dir(): + """Return the test data directory path.""" + return Path("tests/test_data") + + +@pytest.fixture +def require_test_data(test_data_dir): + """Skip test if test data is not available.""" + simple_data_file = test_data_dir / "beam_search_simple_test_data.pkl" + comprehensive_data_file = test_data_dir / "beam_search_comprehensive_test_data.pkl" + + if not simple_data_file.exists() or not comprehensive_data_file.exists(): + pytest.skip("Test data not found. Run scripts/save-data-for-tests.py to generate.") + + return {"simple": simple_data_file, "comprehensive": comprehensive_data_file} + + +@pytest.fixture +def reproducible_seed(): + """Set reproducible seed for tests that need it.""" + + def _set_seed(seed=42): + torch.manual_seed(seed) + np.random.seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + return seed + + return _set_seed + + +# Pytest configuration +def pytest_configure(config): + """Configure pytest for generation tests.""" + config.addinivalue_line("markers", "slow: marks tests as slow (deselect with '-m \"not slow\"')") + config.addinivalue_line("markers", "requires_model: marks tests that require model files") + config.addinivalue_line("markers", "requires_test_data: marks tests that require generated test data") + + +def pytest_collection_modifyitems(config, items): + """Modify test collection to add markers based on fixtures used.""" + for item in items: + # Add requires_model marker if test uses model_files_available fixture + if "model_files_available" in item.fixturenames: + item.add_marker(pytest.mark.requires_model) + + # Add requires_test_data marker if test uses require_test_data fixture + if "require_test_data" in item.fixturenames: + item.add_marker(pytest.mark.requires_test_data) diff --git a/tests/generation/test_batched_beam_search.py b/tests/generation/test_batched_beam_search.py new file mode 100644 index 0000000..41ef2af --- /dev/null +++ b/tests/generation/test_batched_beam_search.py @@ -0,0 +1,323 @@ +""" +Tests for BatchedBeamSearch with variable batch sizes and lengths. +""" + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from directmultistep.generation.tensor_gen import BatchedBeamSearch, BeamSearchOptimized +from directmultistep.utils.dataset import RoutesProcessing + +torch.manual_seed(42) +np.random.seed(42) + + +@pytest.mark.ckptreq +class TestBatchedBeamSearch: + """Test suite for BatchedBeamSearch functionality.""" + + @pytest.fixture(scope="class") + def model_components(self): + """Load model and create batched beam search components.""" + config_path = Path("data/configs/dms_dictionary.yaml") + ckpt_dir = Path("data/checkpoints") + + if not config_path.exists() or not ckpt_dir.exists(): + pytest.skip("Model files not found. Ensure data is downloaded.") + + from directmultistep.generate import load_published_model + + model = load_published_model("flash", ckpt_dir) + rds = RoutesProcessing(metadata_path=config_path) + + device = next(model.parameters()).device + beam_obj = BatchedBeamSearch( + model=model, + beam_size=5, + start_idx=0, + pad_idx=52, + end_idx=22, + max_length=1074, + idx_to_token=rds.idx_to_token, + device=device, + ) + + return model, rds, beam_obj + + def test_batched_beam_search_initialization(self, model_components): + """Test that batched beam search object can be initialized properly.""" + model, rds, beam_obj = model_components + + assert beam_obj.model == model + assert beam_obj.beam_size == 5 + assert beam_obj.start_idx == 0 + assert beam_obj.pad_idx == 52 + assert beam_obj.end_idx == 22 + assert beam_obj.max_length == 1074 + assert isinstance(beam_obj.idx_to_token, dict) + assert isinstance(beam_obj.device, torch.device) + + +@pytest.mark.ckptreq +class TestBatchedVsOptimizedComparison: + """Test that BatchedBeamSearch produces same results as BeamSearchOptimized for single batch.""" + + @pytest.fixture(scope="class") + def model_components(self): + """Load model and create both beam search implementations.""" + config_path = Path("data/configs/dms_dictionary.yaml") + ckpt_dir = Path("data/checkpoints") + + if not config_path.exists() or not ckpt_dir.exists(): + pytest.skip("Model files not found. Ensure data is downloaded.") + + from directmultistep.generate import load_published_model + + model = load_published_model("flash", ckpt_dir) + rds = RoutesProcessing(metadata_path=config_path) + + device = next(model.parameters()).device + + optimized_beam = BeamSearchOptimized( + model=model, + beam_size=5, + start_idx=0, + pad_idx=52, + end_idx=22, + max_length=1074, + idx_to_token=rds.idx_to_token, + device=device, + ) + vec_beam = BatchedBeamSearch( + model=model, + beam_size=5, + start_idx=0, + pad_idx=52, + end_idx=22, + max_length=1074, + idx_to_token=rds.idx_to_token, + device=device, + ) + + return model, rds, optimized_beam, vec_beam + + def test_single_batch_equivalence(self, model_components): + """Test that BatchedBeamSearch gives same results as BeamSearchOptimized for single batch.""" + model, rds, optimized_beam, vec_beam = model_components + + target = "CNCc1ccccc1" + starting_material = None + n_steps = 1 + + from directmultistep.generate import prepare_input_tensors + + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + target, n_steps, starting_material, rds, rds.product_max_length, rds.sm_max_length + ) + + encoder_inp = encoder_inp.to(vec_beam.device) + steps_tens = steps_tens.to(vec_beam.device) if steps_tens is not None else None + path_tens = path_tens.to(vec_beam.device) + + torch.manual_seed(42) + optimized_results = optimized_beam.decode( + src_BC=encoder_inp, + steps_B1=steps_tens, + path_start_BL=path_tens, + progress_bar=True, + ) + + torch.manual_seed(42) + vec_results = vec_beam.decode( + src_BC=encoder_inp, + steps_B1=steps_tens, + path_starts=[path_tens[0]], + progress_bar=True, + ) + + assert len(optimized_results) == 1 + assert len(vec_results) == 1 + + optimized_seqs = [seq for seq, _ in optimized_results[0]] + vec_seqs = [seq for seq, _ in vec_results[0]] + + optimized_probs = [prob for _, prob in optimized_results[0]] + vec_probs = [prob for _, prob in vec_results[0]] + + for i, (o_seq, v_seq) in enumerate(zip(optimized_seqs, vec_seqs, strict=False)): + assert o_seq == v_seq, f"Beam {i}: Sequence mismatch.\nOptimized: {o_seq}\nVectorized: {v_seq}" + + for i, (o_prob, v_prob) in enumerate(zip(optimized_probs, vec_probs, strict=False)): + assert abs(o_prob - v_prob) < 1e-5, ( + f"Beam {i}: Log prob mismatch. Optimized: {o_prob:.6f}, Vectorized: {v_prob:.6f}" + ) + + def test_single_batch_equivalence_with_sm(self, model_components): + """Test equivalence when starting material is provided.""" + model, rds, optimized_beam, vec_beam = model_components + + target = "CNCc1ccccc1" + starting_material = "CN" + n_steps = 1 + + from directmultistep.generate import prepare_input_tensors + + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + target, n_steps, starting_material, rds, rds.product_max_length, rds.sm_max_length + ) + + encoder_inp = encoder_inp.to(vec_beam.device) + steps_tens = steps_tens.to(vec_beam.device) if steps_tens is not None else None + path_tens = path_tens.to(vec_beam.device) + + torch.manual_seed(42) + optimized_results = optimized_beam.decode( + src_BC=encoder_inp, + steps_B1=steps_tens, + path_start_BL=path_tens, + progress_bar=False, + ) + + torch.manual_seed(42) + vec_results = vec_beam.decode( + src_BC=encoder_inp, + steps_B1=steps_tens, + path_starts=[path_tens[0]], + progress_bar=False, + ) + + optimized_seqs = [seq for seq, _ in optimized_results[0]] + vec_seqs = [seq for seq, _ in vec_results[0]] + + optimized_probs = [prob for _, prob in optimized_results[0]] + vec_probs = [prob for _, prob in vec_results[0]] + + for i, (o_seq, v_seq) in enumerate(zip(optimized_seqs, vec_seqs, strict=False)): + assert o_seq == v_seq, f"Beam {i}: Sequence mismatch.\nOptimized: {o_seq}\nVectorized: {v_seq}" + + for i, (o_prob, v_prob) in enumerate(zip(optimized_probs, vec_probs, strict=False)): + assert abs(o_prob - v_prob) < 1e-5, ( + f"Beam {i}: Log prob mismatch. Optimized: {o_prob:.6f}, Vectorized: {v_prob:.6f}" + ) + + def test_multiple_batches_independently_correct(self, model_components): + """Test that batched decoding produces correct results for each batch item independently.""" + model, rds, optimized_beam, vec_beam = model_components + + targets = ["CNCc1ccccc1", "CC(=O)OC1=CC=CC=C1C(=O)O"] + n_steps_list = [1, 1] + + from directmultistep.generate import prepare_batched_input_tensors + + encoder_batch, steps_batch, path_starts, target_lengths = prepare_batched_input_tensors( + targets=targets, + n_steps_list=n_steps_list, + starting_materials=[None, None], + rds=rds, + product_max_length=rds.product_max_length, + sm_max_length=rds.sm_max_length, + ) + + torch.manual_seed(42) + vec_results = vec_beam.decode( + src_BC=encoder_batch.to(vec_beam.device), + steps_B1=steps_batch.to(vec_beam.device) if steps_batch is not None else None, + path_starts=[ps.to(vec_beam.device) for ps in path_starts], + progress_bar=True, + ) + + from directmultistep.generate import prepare_input_tensors + + for idx, (target, n_steps) in enumerate(zip(targets, n_steps_list, strict=False)): + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + target, n_steps, None, rds, rds.product_max_length, rds.sm_max_length + ) + + encoder_inp = encoder_inp.to(optimized_beam.device) + steps_tens = steps_tens.to(optimized_beam.device) if steps_tens is not None else None + path_tens = path_tens.to(optimized_beam.device) + + torch.manual_seed(42) + optimized_results = optimized_beam.decode( + src_BC=encoder_inp, + steps_B1=steps_tens, + path_start_BL=path_tens, + progress_bar=True, + ) + + vec_seqs = [seq for seq, _ in vec_results[idx]] + optimized_seqs = [seq for seq, _ in optimized_results[0]] + + for beam_idx, (v_seq, o_seq) in enumerate(zip(vec_seqs, optimized_seqs, strict=False)): + assert v_seq == o_seq, ( + f"Target {idx} ('{target}'), Beam {beam_idx}: Sequence mismatch.\n" + f"Vectorized: {v_seq}\nOptimized: {o_seq}" + ) + + @pytest.mark.parametrize("beam_size", [5, 20, 50]) + def test_multiple_batches_independently_correct_hard(self, model_components, beam_size): + """Test that batched decoding produces correct results for each batch item independently.""" + model, rds, optimized_beam, vec_beam = model_components + + targets_list = [ + "CC(=O)OC1=CC=CC=C1C(=O)O", + "CNCc1cc(-c2ccccc2F)n(S(=O)(=O)c2cccnc2)c1", + "O=C(c1ccc(NS(=O)(=O)c2cccc3cccnc23)cc1)N1CCN(CC2CC2)CC1", + "COc1ccc(-n2nccn2)c(C(=O)N2CCC[C@@]2(C)c2nc3c(C)c(Cl)ccc3[nH]2)c1", + ] + sms_list = [None, "O=S(=O)(Cl)c1cccnc1", "CCOC(=O)c1ccc(N)cc1", "C[C@@]1(C(=O)O)CCCN1"] + n_steps_list = [1, 2, 5, 4] + + from directmultistep.generate import prepare_batched_input_tensors + + encoder_batch, steps_batch, path_starts, target_lengths = prepare_batched_input_tensors( + targets=targets_list, + n_steps_list=n_steps_list, + starting_materials=sms_list, + rds=rds, + product_max_length=rds.product_max_length, + sm_max_length=rds.sm_max_length, + ) + + torch.manual_seed(42) + vec_beam.beam_size = beam_size + if beam_size != 5: # Only reinitialize if beam size changed from default + vec_beam._init_reusable_tensors() + vec_results = vec_beam.decode( + src_BC=encoder_batch.to(vec_beam.device), + steps_B1=steps_batch.to(vec_beam.device) if steps_batch is not None else None, + path_starts=[ps.to(vec_beam.device) for ps in path_starts], + progress_bar=True, + ) + + from directmultistep.generate import prepare_input_tensors + + for idx, (target, sm, n_steps) in enumerate(zip(targets_list, sms_list, n_steps_list, strict=False)): + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + target, n_steps, sm, rds, rds.product_max_length, rds.sm_max_length + ) + + encoder_inp = encoder_inp.to(optimized_beam.device) + steps_tens = steps_tens.to(optimized_beam.device) if steps_tens is not None else None + path_tens = path_tens.to(optimized_beam.device) + + torch.manual_seed(42) + optimized_beam.beam_size = beam_size + optimized_results = optimized_beam.decode( + src_BC=encoder_inp, + steps_B1=steps_tens, + path_start_BL=path_tens, + progress_bar=True, + ) + + vec_seqs = [seq for seq, _ in vec_results[idx]] + optimized_seqs = [seq for seq, _ in optimized_results[0]] + + for beam_idx, (v_seq, o_seq) in enumerate(zip(vec_seqs, optimized_seqs, strict=False)): + assert v_seq == o_seq, ( + f"Target {idx} ('{target}'), Beam {beam_idx}: Sequence mismatch.\n" + f"Vector: {v_seq}\nOptimized: {o_seq}" + ) diff --git a/tests/generation/test_beam_search.py b/tests/generation/test_beam_search.py new file mode 100644 index 0000000..17cfeac --- /dev/null +++ b/tests/generation/test_beam_search.py @@ -0,0 +1,435 @@ +""" +Tests for beam search functionality in DirectMultiStep. +""" + +import pickle +from pathlib import Path + +import numpy as np +import pytest +import torch + +from directmultistep.generate import create_beam_search, prepare_input_tensors +from directmultistep.utils.dataset import RoutesProcessing + +torch.manual_seed(42) +np.random.seed(42) + + +@pytest.mark.ckptreq +class TestBeamSearch: + """Test suite for beam search functionality.""" + + @pytest.fixture(scope="class") + def test_data(self): + """Load test data generated by save-data-for-tests.py.""" + test_data_path = Path("tests/test_data/beam_search_simple_test_data.pkl") + if not test_data_path.exists(): + pytest.skip("Test data not found. Run scripts/save-data-for-tests.py to generate.") + + with open(test_data_path, "rb") as f: + return pickle.load(f) + + @pytest.fixture(scope="class") + def comprehensive_test_data(self): + """Load comprehensive test data with intermediate results.""" + test_data_path = Path("tests/test_data/beam_search_comprehensive_test_data.pkl") + if not test_data_path.exists(): + pytest.skip("Comprehensive test data not found. Run scripts/save-data-for-tests.py to generate.") + + with open(test_data_path, "rb") as f: + return pickle.load(f) + + @pytest.fixture(scope="class") + def model_components(self): + """Load model and beam search components.""" + config_path = Path("data/configs/dms_dictionary.yaml") + ckpt_dir = Path("data/checkpoints") + + if not config_path.exists() or not ckpt_dir.exists(): + pytest.skip("Model files not found. Ensure data is downloaded.") + + from directmultistep.generate import create_beam_search, load_published_model + + model = load_published_model("flash", ckpt_dir) + rds = RoutesProcessing(metadata_path=config_path) + beam_obj = create_beam_search(model, beam_size=5, rds=rds) + + return model, rds, beam_obj + + def test_beam_search_initialization(self, model_components): + """Test that beam search object can be initialized properly.""" + model, rds, beam_obj = model_components + + assert beam_obj.model == model + assert beam_obj.beam_size == 5 + assert beam_obj.start_idx == 0 + assert beam_obj.pad_idx == 52 + assert beam_obj.end_idx == 22 + assert beam_obj.max_length == 1074 + assert isinstance(beam_obj.idx_to_token, dict) + assert isinstance(beam_obj.device, torch.device) + + def test_beam_search_decode_shape(self, model_components, test_data): + """Test that beam search decode returns correct output shape.""" + model, rds, beam_obj = model_components + + # Use first test case + case_name = "target1" + if case_name not in test_data or test_data[case_name] is None: + pytest.skip(f"Test data for {case_name} not available") + + case_info = test_data[case_name]["case_info"] + + # Prepare input tensors + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + case_info["target"], + case_info["n_steps"], + case_info["starting_material"], + rds, + rds.product_max_length, + rds.sm_max_length, + ) + + # Run beam search + results = beam_obj.decode( + src_BC=encoder_inp.to(beam_obj.device), + steps_B1=steps_tens.to(beam_obj.device) if steps_tens is not None else None, + path_start_BL=path_tens.to(beam_obj.device), + progress_bar=False, + ) + + # Check output shape and structure + assert isinstance(results, list) + assert len(results) == 1 # Single batch + assert isinstance(results[0], list) + assert len(results[0]) == 5 # beam_size + + # Check each beam result + for beam_result in results[0]: + assert isinstance(beam_result, tuple) + assert len(beam_result) == 2 + assert isinstance(beam_result[0], str) # Path string + assert isinstance(beam_result[1], float) # Log probability + + def test_beam_search_reproducibility(self, model_components, comprehensive_test_data): + """Test that beam search produces reproducible results with same seed.""" + model, rds, beam_obj = model_components + + case_name = "target1" + if case_name not in comprehensive_test_data or comprehensive_test_data[case_name] is None: + pytest.skip(f"Comprehensive test data for {case_name} not available") + + case_data = comprehensive_test_data[case_name] + intermediate_data = case_data["intermediate_data"] + + # Set seed for reproducibility + torch.manual_seed(42) + + # Run beam search with same inputs + results = beam_obj.decode( + src_BC=intermediate_data["encoder_input"].to(beam_obj.device), + steps_B1=intermediate_data["steps_tensor"].to(beam_obj.device) + if intermediate_data["steps_tensor"] is not None + else None, + path_start_BL=intermediate_data["path_start_tensor"].to(beam_obj.device), + progress_bar=False, + ) + + # The beam search returns all beam results + assert len(results[0]) == beam_obj.beam_size # Should return full beam results + + # Verify the top result is valid + if results[0]: + actual_top_seq, actual_top_logprob = results[0][0] + assert isinstance(actual_top_seq, str) + assert len(actual_top_seq) > 0 # Should produce some output + assert isinstance(actual_top_logprob, float) + + def test_beam_search_with_different_beam_sizes(self, model_components): + """Test beam search with different beam sizes.""" + model, rds, _ = model_components + + # Test case + target = "CNCc1cc(-c2ccccc2F)n(S(=O)(=O)c2cccnc2)c1" + starting_material = "CN" + n_steps = 2 + + # Prepare input tensors + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + target, n_steps, starting_material, rds, rds.product_max_length, rds.sm_max_length + ) + + device = next(model.parameters()).device + encoder_inp = encoder_inp.to(device) + steps_tens = steps_tens.to(device) if steps_tens is not None else None + path_tens = path_tens.to(device) + + # Test different beam sizes + for beam_size in [1, 3, 5]: + beam_obj = create_beam_search(model, beam_size, rds) + results = beam_obj.decode( + src_BC=encoder_inp, steps_B1=steps_tens, path_start_BL=path_tens, progress_bar=False + ) + + assert len(results[0]) == beam_size + + # Check that all results are valid + for path, log_prob in results[0]: + assert isinstance(path, str) + assert isinstance(log_prob, float) + assert not torch.isnan(torch.tensor(log_prob)) + + def test_beam_search_empty_results(self, model_components): + """Test beam search behavior with invalid inputs.""" + model, rds, beam_obj = model_components + + # Test with characters that are not in the vocabulary + invalid_target = "XYZ123invalid" + + # This should fail gracefully with a KeyError for unknown characters + with pytest.raises(KeyError): + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + invalid_target, 1, None, rds, rds.product_max_length, rds.sm_max_length + ) + + def test_beam_search_edge_cases(self, model_components): + """Test beam search with edge case inputs.""" + model, rds, beam_obj = model_components + + # Test with valid but simple molecule + simple_target = "C" # Methane + + encoder_inp, steps_tens, path_tens = prepare_input_tensors( + simple_target, 1, None, rds, rds.product_max_length, rds.sm_max_length + ) + + results = beam_obj.decode( + src_BC=encoder_inp.to(beam_obj.device), + steps_B1=steps_tens.to(beam_obj.device) if steps_tens is not None else None, + path_start_BL=path_tens.to(beam_obj.device), + progress_bar=False, + ) + + # Should return valid results + assert isinstance(results, list) + assert len(results) == 1 + assert len(results[0]) == beam_obj.beam_size + + def test_beam_search_step_by_step(self, comprehensive_test_data): + """Test individual beam search steps using saved intermediate data.""" + case_name = "target1" + if case_name not in comprehensive_test_data or comprehensive_test_data[case_name] is None: + pytest.skip(f"Comprehensive test data for {case_name} not available") + + case_data = comprehensive_test_data[case_name] + step_data = case_data["beam_search_steps"] + + # Verify that we have step data + assert len(step_data) > 0 + + # Check first step + first_step = step_data[0] + assert "decoder_output" in first_step + assert "log_probs" in first_step + assert "beam_indices" in first_step + assert "step" in first_step + + # Verify tensor shapes are reasonable + decoder_output = first_step["decoder_output"] + log_probs = first_step["log_probs"] + + assert decoder_output.dim() == 2 # Should be (beam_size, vocab_size) + assert log_probs.dim() == 2 # Should be (beam_size, vocab_size) + assert decoder_output.shape == log_probs.shape + + # Check that log probabilities are valid + assert not torch.isnan(log_probs).any() + assert (log_probs <= 0).all() # Log probabilities should be <= 0 + + def test_exact_sequence_reproduction(self, model_components, comprehensive_test_data): + """Test exact reproduction of generated sequences with fixed seed.""" + model, rds, beam_obj = model_components + + case_name = "target1" + if case_name not in comprehensive_test_data or comprehensive_test_data[case_name] is None: + pytest.skip(f"Comprehensive test data for {case_name} not available") + + case_data = comprehensive_test_data[case_name] + intermediate_data = case_data["intermediate_data"] + + # Check if raw_beam_results exists (new format) + if "raw_beam_results" not in case_data: + pytest.skip("Test data needs to be regenerated with raw_beam_results") + + expected_beam_results = case_data["raw_beam_results"] + + # Set seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Call the actual decode method + results = beam_obj.decode( + src_BC=intermediate_data["encoder_input"].to(beam_obj.device), + steps_B1=intermediate_data["steps_tensor"].to(beam_obj.device) + if intermediate_data["steps_tensor"] is not None + else None, + path_start_BL=intermediate_data["path_start_tensor"].to(beam_obj.device), + progress_bar=False, + ) + + # Extract sequences and log probs from results + actual_sequences = [path for path, _ in results[0]] + actual_log_probs = [log_prob for _, log_prob in results[0]] + + # Extract expected sequences and log probs + expected_sequences = [path for path, _ in expected_beam_results[0]] + expected_log_probs = [log_prob for _, log_prob in expected_beam_results[0]] + + # Test exact sequence reproduction + assert len(actual_sequences) == len(expected_sequences), ( + f"Should generate {len(expected_sequences)} sequences, got {len(actual_sequences)}" + ) + + for i, (actual, expected) in enumerate(zip(actual_sequences, expected_sequences, strict=False)): + assert actual == expected, f"Beam {i}: Sequence mismatch.\nExpected: {expected}\nActual: {actual}" + + # Test exact log probability reproduction + for i, (actual_lp, expected_lp) in enumerate(zip(actual_log_probs, expected_log_probs, strict=False)): + assert abs(actual_lp - expected_lp) < 1e-5, ( + f"Beam {i}: Log prob mismatch. Expected: {expected_lp:.6f}, Actual: {actual_lp:.6f}" + ) + + def test_exact_sequence_reproduction_target2(self, model_components, comprehensive_test_data): + """Test exact reproduction for second target molecule.""" + model, rds, beam_obj = model_components + + case_name = "target2" + if case_name not in comprehensive_test_data or comprehensive_test_data[case_name] is None: + pytest.skip(f"Comprehensive test data for {case_name} not available") + + case_data = comprehensive_test_data[case_name] + intermediate_data = case_data["intermediate_data"] + + # Check if raw_beam_results exists (new format) + if "raw_beam_results" not in case_data: + pytest.skip("Test data needs to be regenerated with raw_beam_results") + + expected_beam_results = case_data["raw_beam_results"] + + # Set seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Call the actual decode method + results = beam_obj.decode( + src_BC=intermediate_data["encoder_input"].to(beam_obj.device), + steps_B1=intermediate_data["steps_tensor"].to(beam_obj.device) + if intermediate_data["steps_tensor"] is not None + else None, + path_start_BL=intermediate_data["path_start_tensor"].to(beam_obj.device), + progress_bar=False, + ) + + # Extract sequences and log probs + actual_sequences = [path for path, _ in results[0]] + actual_log_probs = [log_prob for _, log_prob in results[0]] + expected_sequences = [path for path, _ in expected_beam_results[0]] + expected_log_probs = [log_prob for _, log_prob in expected_beam_results[0]] + + # Test exact sequence reproduction + assert len(actual_sequences) == len(expected_sequences) + + for i, (actual, expected) in enumerate(zip(actual_sequences, expected_sequences, strict=False)): + assert actual == expected, f"Target2 Beam {i}: Sequence mismatch.\nExpected: {expected}\nActual: {actual}" + + # Test exact log probability reproduction + for i, (actual_lp, expected_lp) in enumerate(zip(actual_log_probs, expected_log_probs, strict=False)): + assert abs(actual_lp - expected_lp) < 1e-5, ( + f"Target2 Beam {i}: Log prob mismatch. Expected: {expected_lp:.6f}, Actual: {actual_lp:.6f}" + ) + + def test_top_beam_values(self, model_components, comprehensive_test_data): + """Test specific values for the top-ranked beam.""" + model, rds, beam_obj = model_components + + case_name = "target1" + if case_name not in comprehensive_test_data or comprehensive_test_data[case_name] is None: + pytest.skip(f"Comprehensive test data for {case_name} not available") + + case_data = comprehensive_test_data[case_name] + intermediate_data = case_data["intermediate_data"] + + # Check if raw_beam_results exists (new format) + if "raw_beam_results" not in case_data: + pytest.skip("Test data needs to be regenerated with raw_beam_results") + + expected_beam_results = case_data["raw_beam_results"] + + # Set seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Call decode + results = beam_obj.decode( + src_BC=intermediate_data["encoder_input"].to(beam_obj.device), + steps_B1=intermediate_data["steps_tensor"].to(beam_obj.device) + if intermediate_data["steps_tensor"] is not None + else None, + path_start_BL=intermediate_data["path_start_tensor"].to(beam_obj.device), + progress_bar=False, + ) + + # Get top beam (highest scoring) + top_sequence, top_log_prob = results[0][0] + expected_top_sequence, expected_top_log_prob = expected_beam_results[0][0] + + # Assert exact match for top beam + assert top_sequence == expected_top_sequence, ( + f"Top sequence mismatch.\nExpected: {expected_top_sequence}\nActual: {top_sequence}" + ) + + assert abs(top_log_prob - expected_top_log_prob) < 1e-6, ( + f"Top log prob mismatch. Expected: {expected_top_log_prob:.8f}, Actual: {top_log_prob:.8f}" + ) + + # Verify log prob is a reasonable value (negative, not too extreme) + assert top_log_prob < 0, "Log probability should be negative" + assert top_log_prob > -1000, "Log probability should not be extremely negative" + + # Verify sequence is non-empty and contains expected characters + assert len(top_sequence) > 0, "Generated sequence should not be empty" + + def test_beam_ordering(self, model_components, comprehensive_test_data): + """Test that beams are ordered by log probability (highest first).""" + model, rds, beam_obj = model_components + + case_name = "target1" + if case_name not in comprehensive_test_data or comprehensive_test_data[case_name] is None: + pytest.skip(f"Comprehensive test data for {case_name} not available") + + case_data = comprehensive_test_data[case_name] + intermediate_data = case_data["intermediate_data"] + + # Set seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Call decode + results = beam_obj.decode( + src_BC=intermediate_data["encoder_input"].to(beam_obj.device), + steps_B1=intermediate_data["steps_tensor"].to(beam_obj.device) + if intermediate_data["steps_tensor"] is not None + else None, + path_start_BL=intermediate_data["path_start_tensor"].to(beam_obj.device), + progress_bar=False, + ) + + # Extract log probs + log_probs = [log_prob for _, log_prob in results[0]] + + # Verify beams are sorted by log probability (highest to lowest) + for i in range(len(log_probs) - 1): + assert log_probs[i] >= log_probs[i + 1], ( + f"Beams should be ordered by log probability: beam {i} ({log_probs[i]:.6f}) should be >= beam {i + 1} ({log_probs[i + 1]:.6f})" + ) diff --git a/uv.lock b/uv.lock index b764100..90c93a8 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.11" [[package]] @@ -322,7 +322,7 @@ wheels = [ [[package]] name = "directmultistep" -version = "1.1.2" +version = "1.1.3" source = { editable = "." } dependencies = [ { name = "lightning" },