Python module

hf_pipeline

Generalized Token Generation Pipeline

`HFEmbeddingsPipeline`

class max.pipelines.lib.hf_pipeline.HFEmbeddingsPipeline(pipeline_config, torch_device_type)

Generalized token generator pipeline.

Parameters:

pipeline_config (PipelineConfig)
torch_device_type (str)

`encode()`

encode(batch)

Encodes a batch of text inputs.

Parameters:: batch (dict[str, TextContext])
Return type:: dict[str, EmbeddingsOutput]

`prepare_initial_token_inputs()`

prepare_initial_token_inputs(context_batch)

Parameters:: context_batch (list[TextContext])
Return type:: tuple[Tensor, Tensor]

`HFTextGenerationPipeline`

class max.pipelines.lib.hf_pipeline.HFTextGenerationPipeline(pipeline_config, torch_device_type)

HuggingFace text token generator pipeline.

Parameters:

pipeline_config (PipelineConfig)
torch_device_type (str)

`next_token()`

next_token(batch, num_steps)

Provided a batch, process batch inputs, execute the graph for num_steps in a multi-step scenario, then decode the tokens holistically and return the list of decoded tokens.

Parameters:

batch (dict[str, TextContext])
num_steps (int)

Return type:

dict[str, TextGenerationOutput]

`release()`

release(context)

Releases resources associated with this context.

Parameters:: context (TextContext) – Finished context.
Return type:: None

HFEmbeddingsPipeline​

encode()​

prepare_initial_token_inputs()​

HFTextGenerationPipeline​

next_token()​

release()​

`HFEmbeddingsPipeline`

`encode()`

`prepare_initial_token_inputs()`

`HFTextGenerationPipeline`

`next_token()`

`release()`