Fix changed import in infer.py

Signed-off-by: John St. John <jstjohn@nvidia.com>

Author: jstjohn

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/run/infer.py

@@ -59,6 +59,7 @@
 
 import argparse
 import gc
+import inspect
 import json
 import logging
 import os
@@ -91,16 +92,40 @@
     read_run_config,
 )
 from megatron.bridge.utils.common_utils import get_world_size_safe
-from megatron.bridge.utils.instantiate_utils import instantiate
+from megatron.bridge.utils.instantiate_utils import instantiate, register_allowed_target_prefix
 from megatron.core import dist_checkpointing, parallel_state
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.engines.static_engine import StaticInferenceEngine
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
+
+
+try:
+    from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
+        InferenceWrapperConfig,
+    )
+except ImportError:
+
+    @dataclass
+    class InferenceWrapperConfig:
+        """Compatibility shim for MCore versions that removed InferenceWrapperConfig."""
+
+        hidden_size: int
+        inference_max_requests: int
+        inference_max_seq_length: int
+        inference_batch_times_seqlen_threshold: int
+        params_dtype: torch.dtype
+        padded_vocab_size: int
+        nccl_all_reduce_for_prefill: bool = False
+        moe_pad_experts_for_cuda_graph_inference: bool = False
+
+        def add_attributes(self, attributes: dict[str, Any]) -> None:
+            """Match the old MCore config helper used by Evo2TextGenerationController."""
+            for name, value in attributes.items():
+                setattr(self, name, value)
+
+
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.transformer.module import Float16Module
 from megatron.core.utils import get_model_config
@@ -115,6 +140,52 @@
 logger: logging.Logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
+_WRAPPER_INIT_ACCEPTS_CONFIG = (
+    "inference_wrapper_config" in inspect.signature(AbstractModelInferenceWrapper.__init__).parameters
+)
+
+
+class _TextGenerationTokenizerAdapter:
+    """Expose the tokenizer methods expected by MCore's static text-generation path."""
+
+    def __init__(self, tokenizer: _HuggingFaceTokenizer):
+        self._tokenizer = tokenizer
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._tokenizer, name)
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tokenizer.vocab_size
+
+    @property
+    def bos(self) -> Optional[int]:
+        return getattr(self._tokenizer, "bos", None)
+
+    @property
+    def eod(self) -> Optional[int]:
+        return getattr(self._tokenizer, "eod", None)
+
+    def tokenize(self, text: str) -> list[int]:
+        if hasattr(self._tokenizer, "tokenize"):
+            return self._tokenizer.tokenize(text)
+        return self._tokenizer.text_to_ids(text)
+
+    def detokenize(self, tokens: list[int], skip_special_tokens: bool = True) -> str:
+        if hasattr(self._tokenizer, "detokenize"):
+            return self._tokenizer.detokenize(tokens, skip_special_tokens=skip_special_tokens)
+        return self._tokenizer.ids_to_text(tokens)
+
+    def offsets(self, tokens: list[int], text: str) -> list[int]:
+        if hasattr(self._tokenizer, "offsets"):
+            return self._tokenizer.offsets(tokens, text)
+        offsets = []
+        position = 0
+        for token in tokens:
+            offsets.append(position)
+            position += len(self.detokenize([token], skip_special_tokens=False))
+        return offsets
+
 
 # =============================================================================
 # Hardware-Aware Defaults
@@ -243,7 +314,11 @@ def __init__(
             inference_wrapper_config: Configuration with hidden size, vocab size, etc.
             inference_context: Context for managing state and sequence offsets.
         """
-        super().__init__(model, inference_wrapper_config, inference_context)
+        self.inference_wrapper_config = inference_wrapper_config
+        if _WRAPPER_INIT_ACCEPTS_CONFIG:
+            super().__init__(model, inference_wrapper_config, inference_context)
+        else:
+            super().__init__(model, inference_context)
 
     def prep_inference_input(self, prompts_tokens: torch.Tensor) -> Dict[str, Any]:
         """Prepare the inference input data.
@@ -410,6 +485,7 @@ def setup_inference_engine(
         raise FileNotFoundError(f"run_config.yaml not found at {run_config_filename}")
 
     run_config = read_run_config(run_config_filename)
+    register_allowed_target_prefix("bionemo.")
     model_provider = instantiate(run_config["model"])
     logger.info(f"Instantiated model provider: {type(model_provider).__name__}")
 
@@ -446,6 +522,7 @@ def setup_inference_engine(
         tokenizer = _HuggingFaceTokenizer(tokenizer_dir)
     else:
         tokenizer = _HuggingFaceTokenizer(DEFAULT_HF_TOKENIZER_MODEL_PATH)
+    tokenizer = _TextGenerationTokenizerAdapter(tokenizer)
 
     model_provider.vocab_size = tokenizer.vocab_size
     model_provider.should_pad_vocab = True