Skip to main content
Log in

Python module

registry

Model registry, for tracking various model variants.

PipelineRegistry

class max.pipelines.registry.PipelineRegistry(architectures: list[max.pipelines.registry.SupportedArchitecture])

register()

register(architecture: SupportedArchitecture)

Add new architecture to registry.

reset()

reset() → None

retrieve()

retrieve(pipeline_config: PipelineConfig, task: PipelineTask = PipelineTask.TEXT_GENERATION) → tuple[max.pipelines.interfaces.text_generation.PipelineTokenizer, max.pipelines.interfaces.text_generation.TokenGenerator | max.pipelines.interfaces.embeddings_generation.EmbeddingsGenerator]

retrieve_architecture()

retrieve_architecture(model_path: str, trust_remote_code: bool = False) → SupportedArchitecture | None

retrieve_factory()

retrieve_factory(pipeline_config: PipelineConfig, task: PipelineTask = PipelineTask.TEXT_GENERATION) → tuple[max.pipelines.interfaces.text_generation.PipelineTokenizer, Callable[[], Union[max.pipelines.interfaces.text_generation.TokenGenerator, max.pipelines.interfaces.embeddings_generation.EmbeddingsGenerator]]]

validate_pipeline_config()

validate_pipeline_config(pipeline_config: PipelineConfig) → PipelineConfig

SupportedArchitecture

class max.pipelines.registry.SupportedArchitecture(name: str, example_repo_ids: list[str], default_encoding: SupportedEncoding, supported_encodings: dict[max.pipelines.config.SupportedEncoding, list[max.pipelines.kv_cache.cache_params.KVCacheStrategy]], pipeline_model: Type[PipelineModel], task: PipelineTask, tokenizer: Type[TextTokenizer | TextAndVisionTokenizer], default_weights_format: WeightsFormat, rope_type: RopeType = RopeType.none, weight_adapters: dict[max.graph.weights.format.WeightsFormat, Callable[..., dict[str, max.graph.weights.weights.WeightData]]] | None = None)

Initializes a model architecture supported by MAX pipelines.

New architectures should be registered into the PipelineRegistry.

  • Parameters:

    • name – Architecture name.
    • example_repo_ids – Hugging Face repo_id which runs this architecture.
    • default_encoding – Default encoding for the model.
    • supported_encodings – Alternate encodings supported.
    • pipeline_modelPipelineModel class that defines the model graph and execution.
    • task – Which pipeline task should the model run with.
    • tokenizer – Tokenizer used to preprocess model inputs.
    • default_weights_format – The weights format used in pipeline_model.
    • weight_converters – A dictionary of weight loaders to use if the input checkpoint has a different format than the default.

get_pipeline_for_task()

max.pipelines.registry.get_pipeline_for_task(task: PipelineTask, pipeline_config: PipelineConfig) → type[max.pipelines.pipeline.TextGenerationPipeline] | type[max.pipelines.embeddings_pipeline.EmbeddingsPipeline] | type[max.pipelines.speculative_decoding.SpeculativeDecodingTextGenerationPipeline]