TorchModuleTool Reference

TorchModuleTool is the base abstraction for serving PyTorch-backed model inference through Nomad. It defines the preprocess, forward, and postprocess pipeline that turns structured MCP inputs into model calls and model outputs back into structured tool responses.

class nomad.fm_base_tool.TorchModuleTool(*, fm, name=None, description, args_schema, output_schema, batch_size=1, device=<factory>)

Bases: BaseModel, Generic

A helper class for exposing a PyTorch model as an MCP tool for inference. Provides default methods for running the following pipeline:

  1. Preprocess a sequence of Inputs into a ModelInput

  2. Pass ModelInput through the PyTorch model getting ModelOutput

  3. Postprocess ModelOutput into a suitable sequence of Outputs

Complex models (i.e. multi-GPU) may not be fully supported by this class.

Parameters:
fm: torch.nn.Module

The underlying PyTorch model used for inference.

name: str | None

A short name for the foundation model.

description: str

What the foundation model does / how to use it.

args_schema: type[Input]

The input schema for the model.

output_schema: type[Output]

The output schema for the model.

batch_size: int

Inputs to the model will be batched into sets of at most this size.

device: torch.device

The accelerator on which the model is placed.

preprocess(input)

Convert tool input into the form accepted by the model. The input will be of type list[args_schema] with a length of batch_size.

Defaults to torch.data.default_collate.

Parameters:

input (Sequence[Input])

Return type:

ModelInput

_forward(model_inputs)

Process a batch of observations with the model.

Parameters:

model_inputs (ModelInput)

Return type:

ModelOutput

postprocess(model_output)

Postprocess the model’s raw output into a relevant tool output format.

Parameters:

model_output (ModelOutput)

Return type:

Iterable[Output]

classmethod from_pretrained(pretrained_model_name_or_path, **kwargs)

Instantiate tool from a pretrained checkpoint on disk or a remote repo.

Parameters:

pretrained_model_name_or_path (str)

Return type:

TorchModuleTool

model_post_init(_TorchModuleTool__context)

Override this method to perform additional initialization after __init__ and model_construct. This is useful if you want to do some validation that requires the entire model to be initialized.

Return type:

None

final batch(inputs, **kwargs)
Parameters:

inputs (list[Input])

Return type:

list[Output]

final batch_as_completed(inputs, max_concurency=None)
Parameters:

  • inputs (list[Input])

  • max_concurency (int | None)

Return type:

Iterable[Output]

final add_to_fastmcp(server)

Add self as a tool to server.

Parameters:

server (FastMCP)

Return type:

FunctionTool