Skip to content

Commit 0df510f

Browse files
Alias mcore tokenizer methods for eden_megatron static inference
The megatron-core HuggingFaceTokenizer exposes text_to_ids/ids_to_text, but the static engine calls tokenize/detokenize. Apply the same adapter evo2_megatron uses. Signed-off-by: Peter St. John <pstjohn@nvidia.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1 parent efcf501 commit 0df510f

1 file changed

Lines changed: 12 additions & 0 deletions

File tree

  • recipes/eden_megatron/src/bionemo/eden/run

‎recipes/eden_megatron/src/bionemo/eden/run/infer.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,17 @@
109109
logger.setLevel(logging.INFO)
110110

111111

112+
def _adapt_tokenizer_for_generation(tokenizer: Any) -> Any:
113+
"""Alias ``text_to_ids``/``ids_to_text`` as ``tokenize``/``detokenize`` for the mcore static engine."""
114+
if not hasattr(tokenizer, "tokenize") and hasattr(tokenizer, "text_to_ids"):
115+
tokenizer.tokenize = tokenizer.text_to_ids
116+
if not hasattr(tokenizer, "detokenize") and hasattr(tokenizer, "ids_to_text"):
117+
tokenizer.detokenize = tokenizer.ids_to_text
118+
if not hasattr(tokenizer, "bos") and hasattr(tokenizer, "bos_id"):
119+
tokenizer.bos = tokenizer.bos_id
120+
return tokenizer
121+
122+
112123
# =============================================================================
113124
# Hardware-Aware Defaults
114125
# =============================================================================
@@ -420,6 +431,7 @@ def setup_inference_engine(
420431
tokenizer = _HuggingFaceTokenizer(tokenizer_dir)
421432
else:
422433
tokenizer = _HuggingFaceTokenizer(DEFAULT_HF_TOKENIZER_MODEL_PATH)
434+
tokenizer = _adapt_tokenizer_for_generation(tokenizer)
423435

424436
model_provider.vocab_size = tokenizer.vocab_size
425437
model_provider.should_pad_vocab = True

0 commit comments

Comments
 (0)