TextEmbedder

[source]

TextEmbedder class

keras_hub.models.TextEmbedder(*args, compile=True, **kwargs)

Base class for all text embedding tasks.

TextEmbedder tasks wrap a keras_hub.models.Backbone and a keras_hub.models.Preprocessor to create a model that can be used for generating fixed-size sentence embeddings from variable-length text inputs.

All TextEmbedder tasks include a from_preset() constructor which can be used to load a pre-trained config and weights.

Example

# Load a sentence-transformers model.
embedder = keras_hub.models.TextEmbedder.from_preset(
    "all_minilm_l6_v2_en",
)

# Semantic search.
query = "Which planet is known as the Red Planet?"
documents = [
    "Mars is often referred to as the Red Planet.",
    "Venus is often called Earth's twin.",
]
q_emb = embedder.encode_text(query)
d_embs = embedder.encode_text(documents)

[source]

from_preset method

TextEmbedder.from_preset(preset, load_weights=True, **kwargs)

Instantiate a keras_hub.models.Task from a model preset.

A preset is a directory of configs, weights and other file assets used to save and load a pre-trained model. The preset can be passed as one of:

1. a built-in preset identifier like 'bert_base_en'
2. a Kaggle Models handle like 'kaggle://user/bert/keras/bert_base_en'
3. a Hugging Face handle like 'hf://user/bert_base_en'
4. a path to a local preset directory like './bert_base_en'

For any Task subclass, you can run cls.presets.keys() to list all built-in presets available on the class.

This constructor can be called in one of two ways. Either from a task specific base class like keras_hub.models.CausalLM.from_preset(), or from a model class like keras_hub.models.BertTextClassifier.from_preset(). If calling from the a base class, the subclass of the returning object will be inferred from the config in the preset directory.

Arguments

• preset: string. A built-in preset identifier, a Kaggle Models handle, a Hugging Face handle, or a path to a local directory.
• load_weights: bool. If True, saved weights will be loaded into the model architecture. If False, all weights will be randomly initialized.

Examples

# Load a Gemma generative task.
causal_lm = keras_hub.models.CausalLM.from_preset(
    "gemma_2b_en",
)

# Load a Bert classification task.
model = keras_hub.models.TextClassifier.from_preset(
    "bert_base_en",
    num_classes=2,
)

[source]

compile method

TextEmbedder.compile(optimizer="auto", loss=None, metrics=None, **kwargs)

Configures the TextEmbedder task for training.

The TextEmbedder task extends the default compilation signature of keras.Model.compile with a default for optimizer. loss and metrics must be specified by the user, as there is no single standard loss that fits all sentence-transformer training scenarios.

Arguments

• optimizer: "auto", an optimizer name, or a keras.Optimizer instance. Defaults to "auto", which uses the default optimizer for the given model and task. See keras.Model.compile and keras.optimizers for more info on possible optimizer values.
• loss: A loss name or a keras.losses.Loss instance. Must be specified by the user. See keras.Model.compile and keras.losses for more info on possible values.
• metrics: A list of metrics to be evaluated by the model during training and testing. See keras.Model.compile and keras.metrics for more info on possible values.
• **kwargs: See keras.Model.compile for a full list of arguments supported by the compile method.

[source]

save_to_preset method

TextEmbedder.save_to_preset(preset_dir, max_shard_size=10)

Save task to a preset directory.

Arguments

• preset_dir: The path to the local model preset directory.
• max_shard_size: int or float. Maximum size in GB for each sharded file. If None, no sharding will be done. Defaults to 10.

preprocessor property

keras_hub.models.TextEmbedder.preprocessor

A keras_hub.models.Preprocessor layer used to preprocess input.

backbone property

keras_hub.models.TextEmbedder.backbone

A keras_hub.models.Backbone model with the core architecture.