Skip to navigation

nemo_automodel.components.models.beam.tokenization

View as Markdown

Beam tokenizer: render JSON-string tool-call arguments as the dicts Beam’s chat template requires.

Beam’s chat template renders every tool-call argument as its own <rfl:parameter> block, so it needs function.arguments as a mapping and raises on strings. Automodel’s chat datasets store arguments as JSON strings (the OpenAI wire format, and the only form Arrow-backed datasets keep intact). BeamTokenizer loads the regular tokenizer and wraps that instance’s apply_chat_template to parse such strings back into dicts first, which is also what the serving stacks pass to the same template. Other models are unaffected: the wrapper is registered for model_type == "beam" only, and the tokenizer class (and hence save_pretrained output) is unchanged.

Module Contents

Classes

NameDescription
BeamTokenizerTokenizer factory registered for model_type == "beam" in TokenizerRegistry.

Functions

NameDescription
_apply_chat_template_with_dict_arguments-
_message_with_dict_argumentsReturn message with JSON-object string arguments parsed; the input is never mutated.
tool_call_arguments_as_dictsParse JSON-string tool-call arguments into dicts for one conversation or a batch of conversations.
wrap_tokenizerMake tokenizer.apply_chat_template accept JSON-string tool-call arguments.

API

class nemo_automodel.components.models.beam.tokenization.BeamTokenizer()

Tokenizer factory registered for model_type == "beam" in TokenizerRegistry.

nemo_automodel.components.models.beam.tokenization.BeamTokenizer.from_pretrained(
pretrained_model_name_or_path: str | os.PathLike[str],
args: object = (),
kwargs: object = {}
) -> transformers.PreTrainedTokenizerBase
classmethod

Load the tokenizer exactly as the default path does, then wrap apply_chat_template.

nemo_automodel.components.models.beam.tokenization._apply_chat_template_with_dict_arguments(
tokenizer,
conversation,
args = (),
kwargs = {}
)
nemo_automodel.components.models.beam.tokenization._message_with_dict_arguments(
message: typing.Any
) -> typing.Any

Return message with JSON-object string arguments parsed; the input is never mutated.

nemo_automodel.components.models.beam.tokenization.tool_call_arguments_as_dicts(
conversation: typing.Any
) -> typing.Any

Parse JSON-string tool-call arguments into dicts for one conversation or a batch of conversations.

nemo_automodel.components.models.beam.tokenization.wrap_tokenizer(
tokenizer: transformers.PreTrainedTokenizerBase
) -> transformers.PreTrainedTokenizerBase

Make tokenizer.apply_chat_template accept JSON-string tool-call arguments.

The wrapper is an instance attribute, so the tokenizer class (and its save_pretrained output and pickling) stays the one the default loader picked. A functools.partial over a module-level function, not a closure, keeps the tokenizer picklable and deep-copyable.