mp.tasks.text.TextEmbedder

Class that performs embedding extraction on text.

mp.tasks.text.TextEmbedder(
    graph_config: mp.calculators.core.constant_side_packet_calculator_pb2.mediapipe_dot_framework_dot_calculator__pb2.CalculatorGraphConfig
) -> None

This API expects a TFLite model with TFLite Model Metadata that contains the mandatory (described below) input tensors and output tensors. Metadata should contain the input process unit for the model's Tokenizer as well as input / output tensor metadata.

(kTfLiteInt32)

  • 3 input tensors of size [batch_size x bert_max_seq_len] with names "ids", "mask", and "segment_ids" representing the input ids, mask ids, and segment ids respectively.
  • or 1 input tensor of size [batch_size x max_seq_len] representing the input ids.

At least one output tensor with: (kTfLiteFloat32)

  • N components corresponding to the N dimensions of the returned feature vector for this output layer.
  • Either 2 or 4 dimensions, i.e. [1 x N] or [1 x 1 x 1 x N].

graph_config The mediapipe text task graph config proto.

Methods

close

View source

close() -> None

Shuts down the mediapipe text task instance.

Raises
RuntimeError If the mediapipe text task failed to close.

cosine_similarity

View source

@classmethod
cosine_similarity(
    u: mp.tasks.components.containers.Embedding,
    v: mp.tasks.components.containers.Embedding
) -> float

Utility function to compute cosine similarity between two embedding entries.

May return an InvalidArgumentError if e.g. the feature vectors are of different types (quantized vs. float), have different sizes, or have a an L2-norm of 0.

Args
u An embedding entry.
v An embedding entry.

Returns
The cosine similarity for the two embeddings.

Raises
ValueError May return an error if e.g. the feature vectors are of different types (quantized vs. float), have different sizes, or have an L2-norm of 0.

create_from_model_path

View source

@classmethod
create_from_model_path(
    model_path: str
) -> 'TextEmbedder'

Creates an TextEmbedder object from a TensorFlow Lite model and the default TextEmbedderOptions.

Args
model_path Path to the model.

Returns
TextEmbedder object that's created from the model file and the default TextEmbedderOptions.

Raises
ValueError If failed to create TextEmbedder object from the provided file such as invalid file path.
RuntimeError If other types of error occurred.

create_from_options

View source

@classmethod
create_from_options(
    options: mp.tasks.text.TextEmbedderOptions
) -> 'TextEmbedder'

Creates the TextEmbedder object from text embedder options.

Args
options Options for the text embedder task.

Returns
TextEmbedder object that's created from options.

Raises
ValueError If failed to create TextEmbedder object from TextEmbedderOptions such as missing the model.
RuntimeError If other types of error occurred.

embed

View source

embed(
    text: str
) -> mp.tasks.audio.AudioEmbedderResult

Performs text embedding extraction on the provided text.

Args
text The input text.

Returns
An embedding result object that contains a list of embeddings.

Raises
ValueError If any of the input arguments is invalid.
RuntimeError If text embedder failed to run.

__enter__

View source

__enter__()

Returns self upon entering the runtime context.

__exit__

View source

__exit__(
    unused_exc_type, unused_exc_value, unused_traceback
)

Shuts down the mediapipe text task instance on exit of the context manager.

Raises
RuntimeError If the mediapipe text task failed to close.