For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Python module
max.pipelines.architectures.mpnet
MPNet sentence transformer architecture for embeddings generation.
MPNetConfig
class max.pipelines.architectures.mpnet.MPNetConfig(*, dtype, device, pool_embeddings, huggingface_config, max_seq_len)
Bases: ArchConfigWithBoundedMaxSeqLen, ArchConfig
Configuration for MPNet models.
-
Parameters:
device
device: DeviceRef
dtype
dtype: DType
huggingface_config
huggingface_config: AutoConfig
initialize()
classmethod initialize(pipeline_config, model_config=None)
Initializes an MPNetConfig instance from pipeline configuration.
-
Parameters:
-
- pipeline_config (PipelineConfig) – The MAX Engine pipeline configuration.
- model_config (MAXModelConfig | None)
-
Returns:
-
An initialized MPNetConfig instance.
-
Return type:
max_seq_len
max_seq_len: int
pool_embeddings
pool_embeddings: bool
MPNetInputs
class max.pipelines.architectures.mpnet.MPNetInputs(next_tokens_batch, attention_mask, *, kv_cache_inputs=None, lora=None, hidden_states=None)
Bases: ModelInputs
A class representing inputs for the MPNet model.
This class encapsulates the input tensors required for the MPNet model execution:
- next_tokens_batch: A tensor containing the input token IDs
- attention_mask: A tensor containing the extended attention mask
-
Parameters:
-
- next_tokens_batch (Buffer)
- attention_mask (Buffer)
- kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
- lora (LoRAInputs | None)
- hidden_states (Buffer | list[Buffer] | None)
attention_mask
attention_mask: Buffer
next_tokens_batch
next_tokens_batch: Buffer
MPNetPipelineModel
class max.pipelines.architectures.mpnet.MPNetPipelineModel(pipeline_config, session, devices, kv_cache_config, weights, adapter=None, return_logits=ReturnLogits.ALL)
Bases: PipelineModel[TextContext]
-
Parameters:
-
- pipeline_config (PipelineConfig)
- session (InferenceSession)
- devices (list[Device])
- kv_cache_config (KVCacheConfig)
- weights (Weights)
- adapter (WeightsAdapter | None)
- return_logits (ReturnLogits)
execute()
execute(model_inputs)
Executes the graph with the given inputs.
-
Parameters:
-
model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.
-
Returns:
-
ModelOutputs containing the pipeline’s output tensors.
-
Return type:
This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.
load_model()
load_model(session)
-
Parameters:
-
session (InferenceSession)
-
Return type:
model_config_cls
model_config_cls
alias of MPNetConfig
prepare_initial_token_inputs()
prepare_initial_token_inputs(replica_batches, kv_cache_inputs=None, return_n_logits=1)
Prepares the initial inputs to be passed to execute().
The inputs and functionality can vary per model. For example, model
inputs could include encoded tensors, unique IDs per tensor when using
a KV cache manager, and kv_cache_inputs (or None if the model does
not use KV cache). This method typically batches encoded tensors,
claims a KV cache slot if needed, and returns the inputs and caches.
-
Parameters:
-
- replica_batches (Sequence[Sequence[TextContext]])
- kv_cache_inputs (KVCacheInputs[Buffer, Buffer] | None)
- return_n_logits (int)
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!