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.bert
BERT sentence transformer architecture for embeddings generation.
BertInputs
class max.pipelines.architectures.bert.BertInputs(next_tokens_batch: 'Buffer', attention_mask: 'Buffer', *, kv_cache_inputs: 'KVCacheInputs[Buffer, Buffer] | None' = None, lora: 'LoRAInputs | None' = None, hidden_states: 'Buffer | list[Buffer] | None' = None)
Bases: ModelInputs
-
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
BertModelConfig
class max.pipelines.architectures.bert.BertModelConfig(*, dtype, device, pool_embeddings, huggingface_config, max_seq_len)
Bases: ArchConfigWithBoundedMaxSeqLen, ArchConfig
Configuration for Bert models.
-
Parameters:
device
device: DeviceRef
dtype
dtype: DType
huggingface_config
huggingface_config: AutoConfig
initialize()
classmethod initialize(pipeline_config, model_config=None)
Initializes a BertModelConfig instance from pipeline configuration.
-
Parameters:
-
- pipeline_config (PipelineConfig) – The MAX Engine pipeline configuration.
- model_config (MAXModelConfig | None)
-
Returns:
-
An initialized BertModelConfig instance.
-
Return type:
max_seq_len
max_seq_len: int
pool_embeddings
pool_embeddings: bool
BertPipelineModel
class max.pipelines.architectures.bert.BertPipelineModel(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 BertModelConfig
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!