feat: support for batched generation#21
Conversation
WalkthroughAdds a vectorized BatchedBeamSearch implementation, batched-generation helpers and examples, re-exports new APIs, CI and pre-commit tweaks, small logging change, and comprehensive unit tests for single and batched beam search. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Gen as generate_routes_batched
participant Loader as load_published_model
participant Prep as prepare_batched_input_tensors
participant Factory as create_batched_beam_search
participant BBS as BatchedBeamSearch.decode
User->>Gen: call with targets, n_steps_list, starting_materials, beam_size, model/config
alt model provided as name
Gen->>Loader: load_published_model(config_path, ckpt_dir)
Loader-->>Gen: returns torch.nn.Module
else model is module
Gen-->>Gen: use provided module
end
Gen->>Prep: prepare batched src/steps/path_starts/target_lengths
Prep-->>Gen: tensors + metadata
Gen->>Factory: create_batched_beam_search(model, beam_size, rds)
Factory-->>Gen: BatchedBeamSearch instance
Gen->>BBS: decode(src_BC, steps_B1, path_starts, target_lengths, progress_bar)
BBS-->>Gen: batched (sequence, score) results
Gen-->>User: formatted list of route strings
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (6)
examples/batched_generation_example.py (1)
13-14: Consider adding path validation or setup guidance.The hardcoded paths assume data/configs and data/checkpoints exist. Users running this example without setup will encounter unhelpful errors.
Consider adding a check:
+from pathlib import Path + config_path = Path("data/configs/dms_dictionary.yaml") ckpt_dir = Path("data/checkpoints") + +if not config_path.exists() or not ckpt_dir.exists(): + raise FileNotFoundError( + "Required data files not found. " + "Please run scripts/download_files.sh to fetch configs and checkpoints." + )Or add a comment referencing setup instructions.
tests/generation/test_batched_beam_search.py (2)
14-15: Consider removing module-level seeds.Setting seeds at module level can cause test order dependencies and make debugging harder. The tests already set seeds explicitly before each decode call (e.g., lines 125, 133).
The module-level seeds are redundant given the explicit seeding in each test. Consider relying only on the per-test seeds or using the
reproducible_seedfixture from conftest.py.
286-288: Potential test isolation concern with fixture mutation.Mutating
vec_beam.beam_sizeon the class-scoped fixture can cause test order dependencies. After running withbeam_size=50, the fixture retains that value, potentially affecting other tests in the class.Consider either:
- Creating a fresh beam search instance per parametrized test
- Resetting to default beam_size in a cleanup step
- Using function-scoped fixtures for tests that modify beam_size
Similarly, lines 308 mutate
optimized_beam.beam_sizewithout cleanup.tests/generation/conftest.py (1)
11-13: Module-level seeds may cause test order dependencies.Setting seeds at module import conflicts with per-test reproducibility. The
reproducible_seedfixture (lines 87-98) provides controlled seeding. Consider removing module-level seeds.src/directmultistep/__init__.py (2)
17-27: Consider grouping related functions for better API discoverability.The
__all__list is complete and correct. For improved readability and discoverability, consider grouping related functions:__all__ = [ - "BeamSearchOptimized", "BatchedBeamSearch", + "BeamSearchOptimized", + # Factory functions "create_batched_beam_search", "create_beam_search", + # Generation functions "generate_routes", "generate_routes_batched", - "load_published_model", + # Preparation functions "prepare_batched_input_tensors", "prepare_input_tensors", + # Model loading + "load_published_model", ]This groups the API by functionality (factories, generation, preparation, loading) and alphabetizes the classes.
15-15: Consider making logging setup optional or lazy.Calling
setup_logging()at module import time configures logging globally, which can interfere with user applications that have their own logging configuration. Consider:
- Lazy initialization: Only call
setup_logging()when generating routes- Opt-in setup: Provide a function users can call explicitly
- Conditional setup: Check if logging is already configured before calling
Example opt-in approach:
# In __init__.py from directmultistep.utils.logging_config import setup_logging __all__ = [ # ... existing exports ... "setup_logging", # Export for explicit setup ] # Don't call it automatically # setup_logging()Then in the documentation, guide users to call
setup_logging()if they want the default configuration.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (15)
.github/workflows/linting.yml(1 hunks).github/workflows/testing.yml(1 hunks).gitignore(1 hunks).pre-commit-config.yaml(2 hunks)docs/batched-generation.md(1 hunks)examples/batched_generation_example.py(1 hunks)pyproject.toml(2 hunks)scripts/dev/run-beam.py(1 hunks)src/directmultistep/__init__.py(1 hunks)src/directmultistep/generate.py(7 hunks)src/directmultistep/generation/tensor_gen.py(4 hunks)src/directmultistep/model/factory.py(2 hunks)tests/generation/conftest.py(1 hunks)tests/generation/test_batched_beam_search.py(1 hunks)tests/generation/test_beam_search.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
examples/batched_generation_example.py (1)
src/directmultistep/generate.py (1)
generate_routes_batched(266-333)
scripts/dev/run-beam.py (1)
src/directmultistep/generate.py (2)
generate_routes(203-263)generate_routes_batched(266-333)
tests/generation/test_beam_search.py (2)
src/directmultistep/generate.py (3)
create_beam_search(61-75)prepare_input_tensors(95-138)load_published_model(37-58)src/directmultistep/utils/dataset.py (1)
RoutesProcessing(82-155)
src/directmultistep/generate.py (3)
src/directmultistep/generation/tensor_gen.py (3)
BatchedBeamSearch(37-313)decode(71-313)decode(340-475)src/directmultistep/utils/dataset.py (2)
RoutesProcessing(82-155)smile_to_tokens(125-138)src/directmultistep/utils/post_process.py (2)
find_valid_paths(47-76)process_path_single(391-425)
src/directmultistep/generation/tensor_gen.py (1)
src/directmultistep/generation/generation.py (1)
decode(148-235)
tests/generation/test_batched_beam_search.py (3)
src/directmultistep/generation/tensor_gen.py (5)
BatchedBeamSearch(37-313)BeamSearchOptimized(316-475)decode(71-313)decode(340-475)_init_reusable_tensors(64-69)src/directmultistep/utils/dataset.py (1)
RoutesProcessing(82-155)src/directmultistep/generate.py (3)
load_published_model(37-58)prepare_input_tensors(95-138)prepare_batched_input_tensors(141-200)
src/directmultistep/__init__.py (2)
src/directmultistep/generate.py (7)
create_batched_beam_search(78-92)create_beam_search(61-75)generate_routes(203-263)generate_routes_batched(266-333)load_published_model(37-58)prepare_batched_input_tensors(141-200)prepare_input_tensors(95-138)src/directmultistep/generation/tensor_gen.py (2)
BatchedBeamSearch(37-313)BeamSearchOptimized(316-475)
🪛 actionlint (1.7.7)
.github/workflows/linting.yml
22-22: the runner of "actions/setup-python@v4" action is too old to run on GitHub Actions. update the action's version to fix this issue
(action)
.github/workflows/testing.yml
22-22: the runner of "actions/setup-python@v4" action is too old to run on GitHub Actions. update the action's version to fix this issue
(action)
🔇 Additional comments (12)
.gitignore (1)
54-54: Ignore generated collectors artifactsAdding
collectors/keeps generated batch outputs out of version control and aligns with the existing ignore patterns for derived data. Looks good.src/directmultistep/model/factory.py (1)
19-19: LGTM: Clean migration to structured logging.The replacement of
logger.debugis appropriate for parameter count diagnostics. Debug level correctly signals this is development/troubleshooting information rather than user-facing output.Also applies to: 171-171
.github/workflows/testing.yml (1)
3-11: Good addition: Push triggers on main branches.Running tests on push to main/master provides immediate feedback on merged changes and catches integration issues early.
.pre-commit-config.yaml (1)
21-21: LGTM: Consistent with uv dependency management.Invoking mypy via
uv runensures type checking uses the project's managed environment, providing consistent tooling across development and CI..github/workflows/linting.yml (1)
3-11: Good addition: Push triggers on main branches.Consistent with testing.yml, running linting on push provides immediate feedback on code quality.
pyproject.toml (1)
85-86: LGTM: Sensible default for CI environments.Excluding
ckptreqtests by default prevents failures in environments without checkpoint files (e.g., fresh CI runs). Developers can explicitly run these tests withpytest -m ckptreqwhen checkpoints are available.examples/batched_generation_example.py (1)
30-38: LGTM: Correct API usage demonstrating batched generation.The example properly demonstrates
generate_routes_batchedwith:
- Multiple targets with varying complexity
- Mixed starting materials (specified and None)
- Appropriate beam size
This serves as a clear reference for users.
tests/generation/test_batched_beam_search.py (2)
107-156: LGTM: Thorough equivalence testing.The test properly validates that BatchedBeamSearch produces identical results to BeamSearchOptimized by:
- Using deterministic seeds
- Comparing sequences exactly
- Allowing reasonable floating-point tolerance (1e-5) for probabilities
- Providing clear error messages
18-19: LGTM: Clear test requirements and skip logic.The
@pytest.mark.ckptreqdecorator combined with conditional skips in fixtures provides clear feedback when checkpoint files are unavailable. This aligns with the pytest configuration to exclude these tests by default in CI.Also applies to: 64-65
tests/generation/conftest.py (2)
74-83: LGTM: Clear test data requirements with helpful guidance.The
require_test_datafixture provides an actionable skip message directing users to the data generation script. Returning a dict with named keys makes the fixture easy to use in tests.
109-118: LGTM: Automatic marker assignment based on fixtures.The collection hook cleanly auto-tags tests based on their fixture dependencies, making it easy to filter tests by requirement type without manual decorators.
src/directmultistep/__init__.py (1)
3-12: LGTM! Well-organized imports.The imports are logically grouped: high-level API functions from
generateand beam search classes fromtensor_gen. All imported names are properly re-exported in__all__.
| - name: Set up Python 3.11 | ||
| uses: actions/setup-python@v4 | ||
| with: | ||
| python-version: "3.11" |
There was a problem hiding this comment.
Update deprecated GitHub Action version.
Same as testing.yml: actions/setup-python@v4 is outdated. Update to v5 for better Python 3.11+ support.
Apply this diff:
- - name: Set up Python 3.11
- uses: actions/setup-python@v4
- with:
- python-version: "3.11"
+ - name: Set up Python 3.11
+ uses: actions/setup-python@v5
+ with:
+ python-version: "3.11"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| - name: Set up Python 3.11 | |
| uses: actions/setup-python@v4 | |
| with: | |
| python-version: "3.11" | |
| - name: Set up Python 3.11 | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: "3.11" |
🧰 Tools
🪛 actionlint (1.7.7)
22-22: the runner of "actions/setup-python@v4" action is too old to run on GitHub Actions. update the action's version to fix this issue
(action)
🤖 Prompt for AI Agents
.github/workflows/linting.yml lines 21-24: the workflow uses the deprecated
actions/setup-python@v4; update the action reference to actions/setup-python@v5
to ensure current support for Python 3.11+ and replace the version string in the
uses field accordingly.
| - name: Set up Python 3.13 | ||
| uses: actions/setup-python@v4 | ||
| with: | ||
| python-version: "3.13" |
There was a problem hiding this comment.
Update deprecated GitHub Action version.
actions/setup-python@v4 is outdated. The current stable version is v5, which includes improved caching and Python 3.13 support.
Apply this diff:
- - name: Set up Python 3.13
- uses: actions/setup-python@v4
- with:
- python-version: "3.13"
+ - name: Set up Python 3.13
+ uses: actions/setup-python@v5
+ with:
+ python-version: "3.13"🧰 Tools
🪛 actionlint (1.7.7)
22-22: the runner of "actions/setup-python@v4" action is too old to run on GitHub Actions. update the action's version to fix this issue
(action)
🤖 Prompt for AI Agents
In .github/workflows/testing.yml around lines 21 to 24, the workflow uses the
outdated actions/setup-python@v4; update the action reference to the current
stable version by changing the uses line to actions/setup-python@v5 so the job
benefits from improved caching and official Python 3.13 support; ensure no other
inputs change and run CI to verify compatibility.
| torch.stack([s for s in steps_tensors if s is not None]) if all(s is not None for s in steps_tensors) else None | ||
| ) | ||
|
|
There was a problem hiding this comment.
Do not drop per-target n_steps when batching.
Line 197 turns steps_batch into None as soon as any n_steps_list entry is None. That silently discards every other explicit step constraint, so a batch like [2, None, 1] ends up running all three targets without the intended limits. Please keep the provided counts—for example, materialize a (B, 1) tensor and a companion mask, or split into per-item tensors—so that specified values survive all the way to model.encoder. At minimum, fail fast instead of erasing user input.
🤖 Prompt for AI Agents
In src/directmultistep/generate.py around lines 197 to 199, the current
expression collapses steps_batch to None if any entry in n_steps_list is None,
silently discarding explicit per-target step counts; instead preserve per-target
n_steps by materializing a (B,1) tensor (or per-item tensors) plus a companion
mask so entries intentionally unset are represented but other values remain, or
split into per-item tensors before passing to model.encoder; at minimum raise a
clear error if batching with mixed None/int n_steps is unsupported so user input
is not erased.
There was a problem hiding this comment.
@anmorgunov Thank you! Would you like me to open a GitHub issue to track this fix, or do you need any implementation guidance on preserving the per-target step constraints?
| first_step = int(first_steps_B.min().item()) | ||
| max_steps = L - 1 if target_lengths is None else max(target_lengths) | ||
|
|
||
| # Pre-allocate reusable tensors for the loop | ||
| # batch_indices = torch.arange(B, device=self.device).unsqueeze(1).expand(B, S) | ||
| # beam_expand_indices = self.beam_indices.unsqueeze(0).expand(B, -1) | ||
|
|
||
| # 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) | ||
|
|
||
| # Pre-compute beam index patterns for candidate expansion | ||
| beam_idx_pattern = self.beam_indices.unsqueeze(1).expand(-1, S).reshape(-1) | ||
|
|
||
| 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) | ||
|
|
There was a problem hiding this comment.
Respect per-target length caps during decoding.
Line 119 sets a single max_steps, but later iterations never gate beams whose target_lengths[b] has already been reached. As soon as one item requests a longer sequence, every other batch item keeps decoding until that global max, ignoring its own limit. That defeats the “variable target lengths” contract and can yield sequences far longer than the caller asked for. Please carry the per-item limits through the loop—e.g. stash target_lengths as a tensor on the device, clamp max_steps to its max, and break/permanently deactivate beams once step >= target_lengths[b]. A minimal fix looks like:
- max_steps = L - 1 if target_lengths is None else max(target_lengths)
+ if target_lengths is None:
+ target_lengths_B = None
+ max_steps = L - 1
+ else:
+ target_lengths_B = torch.as_tensor(target_lengths, dtype=torch.long, device=self.device)
+ max_steps = min(L - 1, int(target_lengths_B.max().item()))
...
- step_active_B = step >= first_steps_B
+ step_active_B = step >= first_steps_B
+ if target_lengths_B is not None:
+ step_active_B &= step < target_lengths_Band reuse target_lengths_B wherever you need the mask (e.g. batch_finished_B). That keeps shorter targets from over-running their requested length.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| first_step = int(first_steps_B.min().item()) | |
| max_steps = L - 1 if target_lengths is None else max(target_lengths) | |
| # Pre-allocate reusable tensors for the loop | |
| # batch_indices = torch.arange(B, device=self.device).unsqueeze(1).expand(B, S) | |
| # beam_expand_indices = self.beam_indices.unsqueeze(0).expand(B, -1) | |
| # 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) | |
| # Pre-compute beam index patterns for candidate expansion | |
| beam_idx_pattern = self.beam_indices.unsqueeze(1).expand(-1, S).reshape(-1) | |
| 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) | |
| first_step = int(first_steps_B.min().item()) | |
| if target_lengths is None: | |
| target_lengths_B = None | |
| max_steps = L - 1 | |
| else: | |
| target_lengths_B = torch.as_tensor(target_lengths, dtype=torch.long, device=self.device) | |
| max_steps = min(L - 1, int(target_lengths_B.max().item())) | |
| # Pre-allocate reusable tensors for the loop | |
| # batch_indices = torch.arange(B, device=self.device).unsqueeze(1).expand(B, S) | |
| # beam_expand_indices = self.beam_indices.unsqueeze(0).expand(B, -1) | |
| # 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) | |
| # Pre-compute beam index patterns for candidate expansion | |
| beam_idx_pattern = self.beam_indices.unsqueeze(1).expand(-1, S).reshape(-1) | |
| 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 | |
| if target_lengths_B is not None: | |
| step_active_B &= step < target_lengths_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) |
There was a problem hiding this comment.
nah not really a problem, a user doesn't really request target_length, it's just max context window minus prefix length
There was a problem hiding this comment.
@anmorgunov, got it—if target_length is just a computed cap (context window minus prefix) rather than a user requirement, then the global max_steps approach makes sense. As long as sequences naturally terminate via end_idx or the extra decoding is harmless, there's no contract violation. Thanks for clarifying the intent!
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
src/directmultistep/generate.py (1)
178-197: Handle mixedNoneandintvalues inn_steps_list.When
n_steps_listis notNonebut contains someNoneentries (e.g.,[2, None, 1]), line 192 will only append tensors for non-Noneentries, makingsteps_tensorsshorter thanB. Subsequently, line 197'storch.stack(steps_tensors)will create a tensor with incorrect dimensions that doesn't match the batch size, causing a mismatch when passed to the model encoder.Either document that mixed
None/intinn_steps_listis unsupported and raise a clear error, or handle it properly by materializing a(B, 1)tensor with a companion mask to preserve per-target constraints.Apply this diff to add validation:
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 n_steps_list is not None and any(n is None for n in n_steps_list): + raise ValueError("n_steps_list cannot contain None values when specified; use None for the entire list instead") if len(targets) != len(starting_materials):
🧹 Nitpick comments (5)
src/directmultistep/generation/tensor_gen.py (5)
30-34: Remove unusedBeamStatedataclass.The
BeamStatedataclass is defined but never referenced anywhere in this file or the broader codebase. This is dead code.Apply this diff to remove the unused code:
-@dataclass -class BeamState: - sequence: Tensor - score: float - active: bool - -
64-69: Remove unusedbeam_offsetstensor.The
beam_offsetstensor is pre-allocated but never referenced in thedecodemethod. This wastes memory unnecessarily.Apply this diff to remove the unused tensor:
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) - self.beam_offsets = torch.arange(S, device=self.device).unsqueeze(-1) * S
122-123: Remove commented-out code.These commented lines serve no purpose and clutter the implementation.
Apply this diff to remove the dead code:
- # batch_indices = torch.arange(B, device=self.device).unsqueeze(1).expand(B, S) - # beam_expand_indices = self.beam_indices.unsqueeze(0).expand(B, -1) -
131-139: Remove redundantbeam_idx_patterncomputation.
beam_idx_patternis computed identically at lines 131 and 139. The second computation overwrites the first, making it redundant.Apply this diff to remove the redundant line:
- # Pre-compute beam index patterns for candidate expansion - beam_idx_pattern = self.beam_indices.unsqueeze(1).expand(-1, S).reshape(-1) - pbar: Iterable[int] = (
162-170: Simplify redundant active beam checks.Lines 162-163 check
if not active_mask_BS.any(): continue, then lines 169-170 computen_activeand checkif n_active == 0: continue. These checks are redundant sinceactive_mask_flatis just a view ofactive_mask_BS, making the second check unreachable when the first would trigger.Apply this diff to remove the redundant check:
if not active_mask_BS.any(): continue # KEEP ORIGINAL: Only process active sequences active_mask_flat = active_mask_BS.view(B * S) - n_active = active_mask_flat.sum().item() - - if n_active == 0: - continue -
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
.github/workflows/testing.yml(1 hunks)docs/batched-generation.md(1 hunks)src/directmultistep/generate.py(7 hunks)src/directmultistep/generation/tensor_gen.py(4 hunks)
✅ Files skipped from review due to trivial changes (1)
- docs/batched-generation.md
🚧 Files skipped from review as they are similar to previous changes (1)
- .github/workflows/testing.yml
🧰 Additional context used
🧬 Code graph analysis (2)
src/directmultistep/generate.py (3)
src/directmultistep/generation/tensor_gen.py (3)
BatchedBeamSearch(37-308)decode(71-308)decode(335-470)src/directmultistep/utils/dataset.py (2)
RoutesProcessing(82-155)smile_to_tokens(125-138)src/directmultistep/utils/post_process.py (2)
find_valid_paths(47-76)process_path_single(391-425)
src/directmultistep/generation/tensor_gen.py (1)
src/directmultistep/generation/generation.py (1)
decode(148-235)
🔇 Additional comments (10)
src/directmultistep/generation/tensor_gen.py (5)
38-59: LGTM!The constructor properly initializes all required parameters and pre-allocates reusable tensors for efficiency.
61-62: LGTM!Clear and informative
__repr__implementation for debugging.
71-98: LGTM!The method signature is well-designed with proper type hints and sensible defaults. The dtype detection from input and efficient encoder expansion using
contiguous().view()are good practices for supporting both fp32 and fp16.
285-308: LGTM!The result extraction logic correctly handles special tokens (pad/start/end) and properly applies the optional token processor.
357-357: LGTM!The changes to
BeamSearchOptimizedare minor improvements:repeat_interleavefor efficiency, explicit type annotation for the progress bar, and consistent parameter passing. All changes maintain correctness while improving clarity.Also applies to: 376-378, 385-385, 388-388, 414-414
src/directmultistep/generate.py (5)
1-1: LGTM!The new imports (
SequenceandBatchedBeamSearch) are necessary for the batched generation functionality and are properly utilized.Also applies to: 8-8
78-92: LGTM!The
create_batched_beam_searchfunction properly mirrorscreate_beam_search, maintaining consistent parameter initialization for the batched variant.
121-121: LGTM!Removing the individual
unsqueeze(0)calls onprod_tensis a sensible refactor that maintains correct behavior (the batch dimension is still added viaunsqueeze(0)at lines 123/126) while making the tensors easier to stack inprepare_batched_input_tensors.Also applies to: 125-125
212-212: LGTM!Adding
show_progressparameter improves usability by allowing callers to control progress bar display, and it's properly threaded through to the beam search decode call.Also applies to: 249-249
265-333: Approve batched generation function with caveat.The
generate_routes_batchedfunction is well-structured and properly mirrors the single-targetgenerate_routesflow. It correctly validates model constraints per target and processes results individually.However, this function inherits the mixed
None/intissue inn_steps_listfromprepare_batched_input_tensors(flagged separately), which should be addressed.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
src/directmultistep/generation/tensor_gen.py (2)
56-60: Clarify the comment to match implementation.The comment states "Pre-allocate tensors that can be reused" (plural), but only
beam_indicesis allocated. Either update the comment to "Pre-allocate tensor that can be reused" or consider pre-allocating additional tensors if beneficial.
87-88: Consider usingrepeat_interleavefor consistency.Lines 87-88 use
unsqueeze + expand + contiguous + viewfor encoder expansion, while the updatedBeamSearchOptimizedat line 337 usesrepeat_interleave(S, dim=0). The latter is more concise and avoids the explicitcontiguous()call. Consider aligning both implementations for consistency.Apply this diff:
- 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) + enc_src_BSCD = enc_src_BCD.repeat_interleave(S, dim=0) + src_mask_BS11C = src_mask_B11C.repeat_interleave(S, dim=0)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/directmultistep/generation/tensor_gen.py(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
src/directmultistep/generation/tensor_gen.py (1)
src/directmultistep/generation/generation.py (1)
decode(148-235)
🔇 Additional comments (10)
src/directmultistep/generation/tensor_gen.py (10)
90-108: LGTM!Beam initialization and path_starts handling are correctly implemented, with proper dtype propagation for fp16 support.
109-124: LGTM!The
max_stepscalculation and tensor pre-allocation are appropriate. As clarified in previous discussions,target_lengthsrepresents context window limits rather than strict per-target requirements, so the globalmax(target_lengths)approach is valid.The pre-computation of
beam_idx_patternat line 123 is a good optimization that avoids redundant work in the loop.Based on learnings
126-147: LGTM!Active beam filtering and early stopping logic are correctly implemented. The progress bar update provides useful feedback on batch completion.
149-178: LGTM!The optimization to process only active beams is well-implemented. The scatter-gather pattern correctly handles sparse active beams, and dtype propagation ensures fp16 compatibility.
180-208: LGTM!The first-step handling correctly ensures all initial beams derive from the first beam's top-k tokens, which is the expected beam search initialization. The masking logic properly filters invalid candidates.
209-243: LGTM!Inactive beam handling and top-k selection are correctly implemented. Line 214 properly uses
scores_BS.new_full((), float("-inf"))to maintain dtype consistency and avoid implicit upcasting to float32, preserving fp16 benefits throughout.Based on learnings
244-262: LGTM!Sequence updates are correctly implemented. The gather-based reordering efficiently handles beam selection, and the active status update properly tracks completion based on end tokens.
264-288: LGTM!Final result extraction correctly converts beam sequences to strings, handles special tokens appropriately, and applies the optional token processor.
337-337: LGTM!The change to
repeat_interleaveis a cleaner approach for encoder expansion compared to the previousunsqueeze + expand + contiguous + viewpattern.
356-358: LGTM!Adding
dynamic_ncols=Trueimproves progress bar display by adapting to terminal width, enhancing user experience.
1 participant
This PR introduces preliminary support for the batched generation of routes.
See docs/batched-generation.md for detailed intro, but in short, there's now a
generate_routes_batchedthat accepts lists of targets, SMs and n_steps instead of just single elements.Batched generation required reimplementation of beam search, see
BatchedBeamSearchintensor_gen.py. Correctness of the re-implementation was tested against outputs of the existing sequential beam search using DMS Flash model. (runpytest -m "ckptreq"with Flash ckpt indata/checkpoints).========================================================================== test session starts ========================================================================== platform linux -- Python 3.11.4, pytest-8.3.4, pluggy-1.5.0 -- /lambda/nfs/batistalab/DirectMultiStep/.venv/bin/python cachedir: .pytest_cache rootdir: /lambda/nfs/batistalab/DirectMultiStep configfile: pyproject.toml collected 42 items tests/generation/test_batched_beam_search.py::TestBatchedBeamSearch::test_batched_beam_search_initialization PASSED [ 2%] tests/generation/test_batched_beam_search.py::TestBatchedVsOptimizedComparison::test_single_batch_equivalence PASSED [ 4%] tests/generation/test_batched_beam_search.py::TestBatchedVsOptimizedComparison::test_single_batch_equivalence_with_sm PASSED [ 7%] tests/generation/test_batched_beam_search.py::TestBatchedVsOptimizedComparison::test_multiple_batches_independently_correct PASSED [ 9%] tests/generation/test_batched_beam_search.py::TestBatchedVsOptimizedComparison::test_multiple_batches_independently_correct_hard[5] PASSED [ 11%] tests/generation/test_batched_beam_search.py::TestBatchedVsOptimizedComparison::test_multiple_batches_independently_correct_hard[20] PASSED [ 14%] tests/generation/test_batched_beam_search.py::TestBatchedVsOptimizedComparison::test_multiple_batches_independently_correct_hard[50] PASSED [ 16%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_initialization PASSED [ 19%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_decode_shape PASSED [ 21%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_reproducibility PASSED [ 23%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_with_different_beam_sizes PASSED [ 26%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_empty_results PASSED [ 28%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_edge_cases PASSED [ 30%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_search_step_by_step PASSED [ 33%] tests/generation/test_beam_search.py::TestBeamSearch::test_exact_sequence_reproduction PASSED [ 35%] tests/generation/test_beam_search.py::TestBeamSearch::test_exact_sequence_reproduction_target2 PASSED [ 38%] tests/generation/test_beam_search.py::TestBeamSearch::test_top_beam_values PASSED [ 40%] tests/generation/test_beam_search.py::TestBeamSearch::test_beam_ordering PASSED [ 42%] tests/test_preprocess.py::test_filter_mol_nodes[leaves] PASSED [ 45%] tests/test_preprocess.py::test_filter_mol_nodes[depth1] PASSED [ 47%] tests/test_preprocess.py::test_filter_mol_nodes[depth2] PASSED [ 50%] tests/test_preprocess.py::test_filter_mol_nodes[n1_routes_idx0] PASSED [ 52%] tests/test_preprocess.py::test_filter_mol_nodes_invalid_type PASSED [ 54%] tests/test_preprocess.py::test_max_tree_depth[leaves] PASSED [ 57%] tests/test_preprocess.py::test_max_tree_depth[depth1] PASSED [ 59%] tests/test_preprocess.py::test_max_tree_depth[depth2] PASSED [ 61%] tests/test_preprocess.py::test_max_tree_depth[n1_routes_idx0] PASSED [ 64%] tests/test_preprocess.py::test_find_leaves[depth0] PASSED [ 66%] tests/test_preprocess.py::test_find_leaves[depth1] PASSED [ 69%] tests/test_preprocess.py::test_find_leaves[depth2] PASSED [ 71%] tests/test_preprocess.py::test_find_leaves[n1route_idx0] PASSED [ 73%] tests/test_preprocess.py::test_tokenize_smile[data0] PASSED [ 76%] tests/test_preprocess.py::test_tokenize_smile[data1] PASSED [ 78%] tests/test_preprocess.py::test_tokenize_path[data0] PASSED [ 80%] tests/test_preprocess.py::test_generate_permutations_no_children PASSED [ 83%] tests/test_preprocess.py::test_generate_permutations_single_child PASSED [ 85%] tests/test_preprocess.py::test_generate_permutations_multiple_children PASSED [ 88%] tests/test_preprocess.py::test_generate_permutations_nested_children PASSED [ 90%] tests/test_preprocess.py::test_generate_permutations_with_limit PASSED [ 92%] tests/test_preprocess.py::test_generate_permutations_complex_case PASSED [ 95%] tests/test_preprocess.py::test_generate_permutations_parametrized[data0-expected0] PASSED [ 97%] tests/test_preprocess.py::test_generate_permutations_parametrized[data1-expected1] PASSED [100%]A tiny runtime experiment was made on 4 pharma compounds with
n_steps = [1, 2, 5, 4](seescripts/dev/run-beam.py), and the benefits of parallelization are somewhat mixed:At this point, we take it to indicate that careful balancing of targets in the batch is needed to maximize benefits of the new method.
We're going to run a few more internal experiments to gain more confidence into the batched generator before this makes into a v1.2 release
Summary by CodeRabbit