Skip to main content

/

Docs

Nightly

v24.6

Log in

Python module

registry

Model registry, for tracking various model variants.

PipelineRegistry​

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

architecture_details()​

architecture_details(pipeline_config: PipelineConfig) → SupportedArchitecture | None

Return architecture details for pipeline_config if available, None if not found.

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.embeddings_pipeline.EmbeddingsPipeline]

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.embeddings_pipeline.EmbeddingsPipeline]]]

validate_pipeline_config()​

validate_pipeline_config(pipeline_config: PipelineConfig) → PipelineConfig

Update pipeline config with appropriate values if not provided. If invalid config is provided, error out with detailed reason.

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], tokenizer: Type[TextTokenizer | TextAndVisionTokenizer], default_weights_format: WeightsFormat, rope_type: RopeType = RopeType.none, weight_converters: dict[max.pipelines.config.WeightsFormat, Type[max.graph.weights.weights.WeightsConverter]] | 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 – HuggingFace repo_id which runs this architecture.
    • default_encoding – Default encoding for the model.
    • supported_encodings – Alternate encodings supported.
    • pipeline_model – PipelineModel class that defines the model graph and execution.
    • 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.