KerasHub

Model Overview

Sentence Transformers

Sentence Transformers are lightweight BERT-based models from the sentence-transformers project that map sentences and paragraphs into fixed-size dense vector embeddings. These embeddings can be used for semantic textual similarity, semantic search, clustering, paraphrase detection, and information retrieval.

Model Overview

This collection contains KerasHub implementations of popular MiniLM-based Sentence Transformer models, converted from the original sentence-transformers checkpoints. All models use a BERT backbone with mean pooling over token embeddings and produce 384-dimensional output vectors.

The models are organized into four families optimized for different use cases:

Family Use Case Presets
all-MiniLM General-purpose sentence embeddings all_minilm_l6_v2_en, all_minilm_l6_v1_en, all_minilm_l12_v2_en
paraphrase-MiniLM Paraphrase detection & semantic similarity paraphrase_minilm_l3_v2_en, paraphrase_minilm_l6_v2_en, paraphrase_minilm_l12_v2_en
multi-qa-MiniLM Question answering & semantic search multi_qa_minilm_l6_cos_v1_en, multi_qa_minilm_l6_dot_v1_en
msmarco-MiniLM Information retrieval & passage ranking msmarco_minilm_l6_cos_v5_en, msmarco_minilm_l12_cos_v5_en

Key Features

  • Lightweight & Fast: All models use the MiniLM architecture (17M–33M parameters), making them suitable for real-time applications and deployment on resource-constrained hardware.
  • 384-Dimensional Embeddings: All models produce compact 384-dimensional dense vectors, enabling efficient storage and similarity computation.
  • Multi-Backend Support: KerasHub models run on JAX, TensorFlow, and PyTorch backends.
  • Mean Pooling with Attention Mask: Embeddings are computed using attention-aware mean pooling over the BERT token outputs, correctly ignoring padding tokens.
  • L2 Normalization (select models): Models in the all-*, multi-qa-*-cos, and msmarco-*-cos families produce L2-normalized embeddings, enabling fast cosine similarity via dot product.

Links

Installation

Keras and KerasHub can be installed with:

pip install -U -q keras-hub
pip install -U -q keras

Jax, TensorFlow, and Torch come preinstalled in Kaggle Notebooks. For instructions on installing them in another environment see the Keras Getting Started page.

Example Usage

Semantic Search

Use encode_text() for queries and encode_documents() for documents, then compute similarity:

import keras_hub
# Load a pre-trained sentence embedder.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
   "multi_qa_minilm_l6_dot_v1_en",
)
# Encode a query and a set of documents.
query = "How do I get a replacement Medicare card?"
documents = [
    "Medicare replacement cards can be requested online or by phone.",
    "The longest river in the world is the Nile.",
    "You can apply for a new Medicare card at ssa.gov.",
]
query_embedding = embedder.encode_text(query)
document_embeddings = embedder.encode_documents(documents)
# Compute cosine similarity (embeddings are L2-normalized by default).
similarities = embedder.similarity(query_embedding, document_embeddings)
print(similarities)
# Best match:
best_idx = int(similarities[0].argmax())
print(f"Best match: {documents[best_idx]}")

Using the Base Class

You can also use the TextEmbedder base class to load any sentence-transformer preset:

import keras_hub

embedder = keras_hub.models.TextEmbedder.from_preset(
   "multi_qa_minilm_l6_dot_v1_en",
)

# Encode a batch of sentences.
sentences = [
    "That is a happy person",
    "That is a very happy person",
    "Today is a sunny day",
]
embeddings = embedder.predict(sentences)
print(embeddings.shape)  # (3, 384)

Generating Embeddings for Clustering

import keras_hub
import numpy as np

embedder = keras_hub.models.BertTextEmbedder.from_preset(
    "multi_qa_minilm_l6_dot_v1_en",
)

# Two groups of semantically related sentences.
group_a = [
    "The cat sits on the mat.",
    "A kitten is resting on a rug.",
]
group_b = [
    "The stock market rallied today.",
    "Financial markets saw significant gains.",
]

all_sentences = group_a + group_b
embeddings = embedder.encode_text(all_sentences)
print(embeddings.shape)  # (4, 384)

# Compute pairwise cosine similarity across all sentences.
# This produces a (4, 4) matrix where entry [i][j] is the
# similarity between sentence i and sentence j.
similarity_matrix = embedder.similarity(embeddings, embeddings)
print(similarity_matrix)
# Sentences 0-1 (cat/kitten) will have high mutual similarity.
# Sentences 2-3 (markets) will have high mutual similarity.
# Cross-group similarity (e.g., 0 vs 2) will be low.

Preprocessed Input (No Preprocessing)

If you want to handle tokenization yourself:


import keras_hub
import numpy as np

# Load without preprocessor.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
   "multi_qa_minilm_l6_dot_v1_en",
    preprocessor=None,
)

# Provide pre-tokenized inputs.
features = {
    "token_ids": np.ones(shape=(2, 12), dtype="int32"),
    "segment_ids": np.array(
        [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]] * 2
    ),
    "padding_mask": np.array(
        [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2
    ),
}
embeddings = embedder.predict(features)
print(embeddings.shape)  # (2, 384)

Choosing a Pooling Mode

By default, models use mean pooling with L2 normalization. You can customize this:

import keras_hub

# CLS pooling (use [CLS] token representation).
embedder = keras_hub.models.BertTextEmbedder.from_preset(
   "multi_qa_minilm_l6_dot_v1_en",
    pooling_mode="cls",
)

# Max pooling.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
    "multi_qa_minilm_l6_dot_v1_en",
    pooling_mode="max",
)

# Disable L2 normalization.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
    "multi_qa_minilm_l6_dot_v1_en",
    normalize=False,
)

Example Usage with Hugging Face URI

Semantic Search

Use encode_text() for queries and encode_documents() for documents, then compute similarity:

import keras_hub
# Load a pre-trained sentence embedder.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
   "hf://keras/multi_qa_minilm_l6_dot_v1_en",
)
# Encode a query and a set of documents.
query = "How do I get a replacement Medicare card?"
documents = [
    "Medicare replacement cards can be requested online or by phone.",
    "The longest river in the world is the Nile.",
    "You can apply for a new Medicare card at ssa.gov.",
]
query_embedding = embedder.encode_text(query)
document_embeddings = embedder.encode_documents(documents)
# Compute cosine similarity (embeddings are L2-normalized by default).
similarities = embedder.similarity(query_embedding, document_embeddings)
print(similarities)
# Best match:
best_idx = int(similarities[0].argmax())
print(f"Best match: {documents[best_idx]}")

Using the Base Class

You can also use the TextEmbedder base class to load any sentence-transformer preset:

import keras_hub

embedder = keras_hub.models.TextEmbedder.from_preset(
   "hf://keras/multi_qa_minilm_l6_dot_v1_en",
)

# Encode a batch of sentences.
sentences = [
    "That is a happy person",
    "That is a very happy person",
    "Today is a sunny day",
]
embeddings = embedder.predict(sentences)
print(embeddings.shape)  # (3, 384)

Generating Embeddings for Clustering

import keras_hub
import numpy as np

embedder = keras_hub.models.BertTextEmbedder.from_preset(
    "hf://keras/multi_qa_minilm_l6_dot_v1_en",
)

# Two groups of semantically related sentences.
group_a = [
    "The cat sits on the mat.",
    "A kitten is resting on a rug.",
]
group_b = [
    "The stock market rallied today.",
    "Financial markets saw significant gains.",
]

all_sentences = group_a + group_b
embeddings = embedder.encode_text(all_sentences)
print(embeddings.shape)  # (4, 384)

# Compute pairwise cosine similarity across all sentences.
# This produces a (4, 4) matrix where entry [i][j] is the
# similarity between sentence i and sentence j.
similarity_matrix = embedder.similarity(embeddings, embeddings)
print(similarity_matrix)
# Sentences 0-1 (cat/kitten) will have high mutual similarity.
# Sentences 2-3 (markets) will have high mutual similarity.
# Cross-group similarity (e.g., 0 vs 2) will be low.

Preprocessed Input (No Preprocessing)

If you want to handle tokenization yourself:


import keras_hub
import numpy as np

# Load without preprocessor.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
   "hf://keras/multi_qa_minilm_l6_dot_v1_en",
    preprocessor=None,
)

# Provide pre-tokenized inputs.
features = {
    "token_ids": np.ones(shape=(2, 12), dtype="int32"),
    "segment_ids": np.array(
        [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]] * 2
    ),
    "padding_mask": np.array(
        [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2
    ),
}
embeddings = embedder.predict(features)
print(embeddings.shape)  # (2, 384)

Choosing a Pooling Mode

By default, models use mean pooling with L2 normalization. You can customize this:

import keras_hub

# CLS pooling (use [CLS] token representation).
embedder = keras_hub.models.BertTextEmbedder.from_preset(
   "hf://keras/multi_qa_minilm_l6_dot_v1_en",
    pooling_mode="cls",
)

# Max pooling.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
    "hf://keras/multi_qa_minilm_l6_dot_v1_en",
    pooling_mode="max",
)

# Disable L2 normalization.
embedder = keras_hub.models.BertTextEmbedder.from_preset(
    "hf://keras/multi_qa_minilm_l6_dot_v1_en",
    normalize=False,
)
Downloads last month
24
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Collection including keras/multi_qa_minilm_l6_dot_v1_en

Paper for keras/multi_qa_minilm_l6_dot_v1_en