pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")

queries = ["Which planet is known as the Red Planet?"]
documents = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
]

# encode_query and encode_document apply the "query" and "document" task prompts
query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings.shape)
# (1, 768) (2, 768)

# (1, 2) matrix of cosine similarities
print(model.similarity(query_embeddings, document_embeddings))

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer


model = AutoModel.from_pretrained("google/embeddinggemma-2", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("google/embeddinggemma-2")

# Prepend the task prompt directly to plain text inputs (see Task prompts below)
sentences = [
    "task: search result | query: Which planet is known as the Red Planet?",
    "title: none | text: Venus is often called Earth's twin because of its similar size and proximity.",
    "title: none | text: Mars, known for its reddish appearance, is often referred to as the Red Planet.",
]
inputs = tokenizer(sentences, padding=True, return_tensors="pt").to(model.device)

with torch.no_grad():
    # (batch_size, sequence_length, config.text_config.embedding_dim)
    token_embeddings = model(**inputs).last_hidden_state

# Mask-aware mean pooling over non-padded tokens, then L2 normalization in float32
mask = inputs["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
sentence_embeddings = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
sentence_embeddings = F.normalize(sentence_embeddings.float(), p=2, dim=-1)

query_embedding, document_embeddings = sentence_embeddings[:1], sentence_embeddings[1:]
print(query_embedding @ document_embeddings.T)

import torch.nn.functional as F
from sentence_transformers import SentenceTransformer


# truncate_dim can also be set once at load time: SentenceTransformer(..., truncate_dim=256)
model = SentenceTransformer("google/embeddinggemma-2")

queries = ["Which planet is known as the Red Planet?"]
documents = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
]

query_embeddings = model.encode_query(queries, truncate_dim=256, convert_to_tensor=True)
document_embeddings = model.encode_document(documents, truncate_dim=256, convert_to_tensor=True)

query_embeddings = F.normalize(query_embeddings.float(), p=2, dim=-1)
document_embeddings = F.normalize(document_embeddings.float(), p=2, dim=-1)

# (1, 256) and (2, 256): a 3x smaller index than the full 768 dimensions
print(query_embeddings.shape, document_embeddings.shape)

# the Mars document still ranks first
print(model.similarity(query_embeddings, document_embeddings))

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer


model = AutoModel.from_pretrained("google/embeddinggemma-2", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("google/embeddinggemma-2")

sentences = [
    "task: search result | query: Which planet is known as the Red Planet?",
    "title: none | text: Venus is often called Earth's twin because of its similar size and proximity.",
    "title: none | text: Mars, known for its reddish appearance, is often referred to as the Red Planet.",
]
inputs = tokenizer(sentences, padding=True, return_tensors="pt").to(model.device)

with torch.no_grad():
    token_embeddings = model(**inputs).last_hidden_state

mask = inputs["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
sentence_embeddings = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)

# Slice to the first 256 dimensions, then normalize
sentence_embeddings = F.normalize(sentence_embeddings[:, :256].float(), p=2, dim=-1)
print(sentence_embeddings.shape)
# torch.Size([3, 256])

query_embedding, document_embeddings = sentence_embeddings[:1], sentence_embeddings[1:]

# the Mars document still ranks first
print(query_embedding @ document_embeddings.T)

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")
IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

# Embed each modality on its own into the shared (768,) space
candidates = model.encode(
    [
        {"text": "A photo of a cat"},
        {"image": IMAGE},
        {"audio": "path/to/audio.wav"},
        {"video": "path/to/video.mp4"},
    ]
)
print(candidates.shape)
# (4, 768)

# Cross-modal retrieval: rank text, image, audio, and video candidates against a text query
query = model.encode_query(["A fluffy cat outdoors"])
print(model.similarity(query, candidates))
# tensor([[...]]) of shape (1, 4)

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoProcessor


model = AutoModel.from_pretrained("google/embeddinggemma-2", device_map="auto")
processor = AutoProcessor.from_pretrained("google/embeddinggemma-2")

IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

conversations = [
    [{"role": "user", "content": [{"type": "text", "text": "A photo of a cat"}]}],
    [{"role": "user", "content": [{"type": "image", "url": IMAGE}]}],
    [{"role": "user", "content": [{"type": "audio", "url": "path/to/audio.wav"}]}],
    [{"role": "user", "content": [{"type": "video", "url": "path/to/video.mp4"}]}],
]

inputs = processor.apply_chat_template(
    conversations,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    token_embeddings = model(**inputs).last_hidden_state

mask = inputs["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
embeddings = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
embeddings = F.normalize(embeddings.float(), p=2, dim=-1)
print(embeddings.shape)
# torch.Size([4, 768])

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")
IMAGE_1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
IMAGE_2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png"

composed_embeddings = model.encode(
    [
        # text + image
        {"text": "A photo of a cat", "image": IMAGE_1},
        # text + audio
        {"text": "A song", "audio": "path/to/audio.wav"},
        # image + audio (no text required)
        {"image": IMAGE_1, "audio": "path/to/audio.wav"},
        # multiple images + audio + text in one embedding
        {"image": [IMAGE_1, IMAGE_2], "audio": "path/to/audio.wav", "text": "Compare both cats"},
    ]
)
print(composed_embeddings.shape)
# (4, 768)

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoProcessor


model = AutoModel.from_pretrained("google/embeddinggemma-2", device_map="auto")
processor = AutoProcessor.from_pretrained("google/embeddinggemma-2")

IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

conversations = [
    # text + image
    [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "A photo of a cat"},
                {"type": "image", "url": IMAGE},
            ],
        }
    ],
    # text + audio
    [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "A song"},
                {"type": "audio", "url": "path/to/audio.wav"},
            ],
        }
    ],
    # image + audio (no text required)
    [
        {
            "role": "user",
            "content": [
                {"type": "image", "url": IMAGE},
                {"type": "audio", "url": "path/to/audio.wav"},
            ],
        }
    ],
]

inputs = processor.apply_chat_template(
    conversations,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    token_embeddings = model(**inputs).last_hidden_state

mask = inputs["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
embeddings = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
embeddings = F.normalize(embeddings.float(), p=2, dim=-1)
print(embeddings.shape)
# torch.Size([3, 768])

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")
IMAGE_1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
IMAGE_2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png"

# 1. Automatic ordering: follows the dictionary's key order
#    -> <bos> [prompt] <image_tokens> <text_tokens> <eos>
image_first = model.encode({"image": IMAGE_1, "text": "A photo of a cat"}, prompt_name="document")

#    -> <bos> [prompt] <text_tokens> <image_tokens> <eos>
text_first = model.encode({"text": "A photo of a cat", "image": IMAGE_1}, prompt_name="document")

# 2. Manual placeholders: interleave media at exact positions in the text
#    No extra placeholders are inserted; counts must match the passed media inputs.
interleaved = model.encode(
    {
        "text": "A jacket similar to <|image|> or <|image|> featured in <|audio|>",
        "image": [IMAGE_1, IMAGE_2],
        "audio": "path/to/audio.wav",
    },
    prompt_name="query",
)

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoProcessor


model = AutoModel.from_pretrained("google/embeddinggemma-2", device_map="auto")
processor = AutoProcessor.from_pretrained("google/embeddinggemma-2")

IMAGE_1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
IMAGE_2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png"

# 1. Automatic ordering: follows the content list order after any system prompt
image_first_messages = [
    {"role": "system", "content": "title: none | text: "},
    {
        "role": "user",
        "content": [
            {"type": "image", "url": IMAGE_1},
            {"type": "text", "text": "A photo of a cat"},
        ],
    },
]

# 2. Manual placeholders: interleave media at exact positions in the text
interleaved_messages = [
    {"role": "system", "content": "task: search result | query: "},
    {
        "role": "user",
        "content": [
            {"type": "image", "url": IMAGE_1},
            {"type": "image", "url": IMAGE_2},
            {"type": "audio", "url": "path/to/audio.wav"},
            {"type": "text", "text": "A jacket similar to <|image|> or <|image|> featured in <|audio|>"},
        ],
    },
]

inputs = processor.apply_chat_template(
    [image_first_messages, interleaved_messages],
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    token_embeddings = model(**inputs).last_hidden_state

mask = inputs["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
embeddings = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
embeddings = F.normalize(embeddings.float(), p=2, dim=-1)

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")
IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

embeddings = model.encode(
    [
        "A photo of a cat",
        {"image": IMAGE},
        {"text": "A photo of a cat", "image": IMAGE},
        {"audio": "path/to/audio.wav"},
    ],
    prompt_name="document",
)
print(embeddings.shape)
# (4, 768)

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoProcessor


model = AutoModel.from_pretrained("google/embeddinggemma-2", device_map="auto")
processor = AutoProcessor.from_pretrained("google/embeddinggemma-2")

IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

conversations = [
    [
        {"role": "system", "content": "title: none | text: "},
        {"role": "user", "content": [{"type": "text", "text": "A photo of a cat"}]},
    ],
    [
        {"role": "system", "content": "title: none | text: "},
        {"role": "user", "content": [{"type": "image", "url": IMAGE}]},
    ],
    [
        {"role": "system", "content": "title: none | text: "},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "A photo of a cat"},
                {"type": "image", "url": IMAGE},
            ],
        },
    ],
    [
        {"role": "system", "content": "title: none | text: "},
        {"role": "user", "content": [{"type": "audio", "url": "path/to/audio.wav"}]},
    ],
]

inputs = processor.apply_chat_template(
    conversations,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    token_embeddings = model(**inputs).last_hidden_state

mask = inputs["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype)
embeddings = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
embeddings = F.normalize(embeddings.float(), p=2, dim=-1)
print(embeddings.shape)
# torch.Size([4, 768])

from sentence_transformers import SentenceTransformer


# Text-only deployment (skips both vision and audio towers)
text_model = SentenceTransformer(
    "google/embeddinggemma-2",
    config_kwargs={"vision_config": None, "audio_config": None},
)

# Vision + text deployment (skips audio tower)
vision_text_model = SentenceTransformer(
    "google/embeddinggemma-2",
    config_kwargs={"audio_config": None},
)

from transformers import AutoModel


# Text-only deployment (skips both vision and audio towers)
text_model = AutoModel.from_pretrained(
    "google/embeddinggemma-2",
    vision_config=None,
    audio_config=None,
    device_map="auto",
)

# Vision + text deployment (skips audio tower)
vision_text_model = AutoModel.from_pretrained(
    "google/embeddinggemma-2",
    audio_config=None,
    device_map="auto",
)

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")
IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

image_embedding = model.encode(
    {"image": IMAGE},
    processing_kwargs={"image": {"max_soft_tokens": 70}},
)

from transformers import AutoProcessor


processor = AutoProcessor.from_pretrained("google/embeddinggemma-2")
IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"

inputs = processor(images=[IMAGE], max_soft_tokens=70, return_tensors="pt")
print(inputs["input_ids"].shape, inputs["pixel_values"].shape)
# torch.Size([1, 67]) torch.Size([1, 630, 768])

from sentence_transformers import SentenceTransformer


model = SentenceTransformer("google/embeddinggemma-2")

video_embedding = model.encode(
    {"video": "path/to/video.mp4"},
    processing_kwargs={"video": {"add_timestamps": True, "fps": 2, "max_frames": 16, "overflow_strategy": "truncate"}},
)

from transformers import AutoProcessor


processor = AutoProcessor.from_pretrained("google/embeddinggemma-2")

inputs = processor(
    videos=["path/to/video.mp4"],
    add_timestamps=True,
    fps=2,
    max_frames=16,
    overflow_strategy="truncate",
    return_tensors="pt",
)