Attention: This MediaPipe Solutions Preview is an early release. Learn more.

mp.tasks.text.TextEmbedder

Stay organized with collections Save and categorize content based on your preferences.

Class that performs embedding extraction on text.

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

Shuts down the mediapipe text task instance.

Raises
RuntimeError If the mediapipe text task failed to close.

cosine_similarity

View source

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

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

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

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

Returns self upon entering the runtime context.

__exit__

View source

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

Raises
RuntimeError If the mediapipe text task failed to close.