Multimodal Search with Gemini Embedding 2 in Haystack

# pip install haystack-ai google-genai-haystack
from datasets import load_dataset
from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.embedders.google_genai import (
    GoogleGenAIDocumentEmbedder, GoogleGenAITextEmbedder
)

document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")

docs = [
    Document(content="The capybara is the largest rodent in the world and is native to South America, where it lives near rivers, lakes, and wetlands. It is highly social and often seen relaxing in groups, spending much of its time swimming or soaking in water. Capybaras communicate through whistles, barks, and purr-like sounds."),
    Document(content="Dogs are domesticated mammals known for their loyalty, intelligence, and strong bond with humans. They have been bred for thousands of years for roles such as companionship, hunting, guarding, and assisting people with various tasks. Different breeds vary widely in size, temperament, and abilities."),
    Document(content="The tiger is the largest species of big cat and is recognized by its distinctive orange coat with black stripes. It is a powerful solitary predator that inhabits forests, grasslands, and wetlands across parts of Asia. Tigers are excellent swimmers and rely on stealth and strength to hunt prey."),
    Document(content="The giraffe is the tallest land animal on Earth, easily identified by its long neck and distinctive spotted coat. It uses its height to reach leaves high in acacia trees and roams the savannas and open woodlands of Africa. Despite its long neck, a giraffe has the same number of neck vertebrae as most mammals."),
    Document(content="Elephants are the largest land animals and are known for their intelligence, strong family bonds, and remarkable memory. They use their trunks for breathing, grasping objects, and communication. Elephants live in complex social groups led by a matriarch."),
    Document(content="Penguins are flightless birds that live primarily in the Southern Hemisphere, especially in Antarctica. They are excellent swimmers, using their flipper-like wings to move through the water while hunting fish, squid, and krill."),
    Document(content="Dolphins are highly intelligent marine mammals known for their playful behavior and complex communication. They live in social groups called pods and use echolocation to navigate and locate prey in the ocean."),
    Document(content="Owls are nocturnal birds of prey with excellent night vision and silent flight. They hunt small mammals, insects, and other birds, relying on their sharp talons and keen hearing to detect prey in darkness."),
    Document(content="Red pandas are small mammals native to the eastern Himalayas and southwestern China. They have reddish-brown fur, bushy tails, and spend most of their time in trees. Their diet mainly consists of bamboo, though they may also eat fruits and insects."),
    Document(content="Octopuses are highly intelligent sea creatures with eight arms lined with suckers. They are known for their ability to change color and texture to camouflage with their surroundings, helping them avoid predators and ambush prey."),
]

doc_embedder = GoogleGenAIDocumentEmbedder(
    model="gemini-embedding-2-preview", 
    batch_size=5, 
    config={
		    "task_type": "RETRIEVAL_DOCUMENT",
		    "output_dimensionality": 768 
		}
)
docs_with_embeddings = doc_embedder.run(docs)
document_store.write_documents(docs_with_embeddings["documents"])

from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack_integrations.components.embedders.google_genai import GoogleGenAITextEmbedder

embedding_retriever = InMemoryEmbeddingRetriever(document_store=document_store)

query = "animal that communicates with whistles and barks"
text_embedder = GoogleGenAITextEmbedder(
    model="gemini-embedding-2-preview", 
    config={"task_type": "RETRIEVAL_QUERY"}
)
query_embedding = text_embedder.run(query)["embedding"]

result = embedding_retriever.run(query_embedding=query_embedding, top_k=2)

for doc in result["documents"]:
    print(doc.meta)
    print(doc.content)
    print(doc.score)
    print("-" * 10)

from haystack_integrations.components.embedders.google_genai import GoogleGenAIMultimodalDocumentEmbedder
from haystack import Document

docs = [
    Document(meta={"file_path": "kangaroo.mp4"}),
    Document(meta={"file_path": "tiger.jpg"}),
    Document(meta={"file_path": "sample.pdf"}),
    Document(meta={"file_path": "dog.jpg"}),
    Document(meta={"file_path": "cat.jpg"})
]

doc_multimodal_embedder = GoogleGenAIMultimodalDocumentEmbedder(
    model="gemini-embedding-2-preview", 
    config={"task_type": "RETRIEVAL_DOCUMENT"}
)
docs_with_embeddings = doc_multimodal_embedder.run(docs)
document_store.write_documents(docs_with_embeddings["documents"])

image_doc = Document(meta={"file_path": "dog.jpg"})
image_embedder = GoogleGenAIMultimodalDocumentEmbedder(
    model="gemini-embedding-2-preview", 
    config={"task_type": "RETRIEVAL_QUERY"}
)

# Create the embedding for the image
image_embedding = image_embedder.run([image_doc])["documents"][0].embedding

# Find the most semantically similar documents in the vector database
results = embedding_retriever.run(query_embedding=image_embedding, top_k=3)["documents"]

for doc in results:
    print(doc.content)
    print(doc.score)
    print("-" * 100)