NetraEmbed is a multilingual multimodal embedding model that encodes both visual documents and text queries into single dense vectors. It supports multiple languages and enables efficient similarity search at multiple embedding dimensions (768, 1536, 2560) through Matryoshka representation learning.
1import torch
2from PIL import Image
3from colpali_engine.models import BiGemma3, BiGemmaProcessor3
4
5# Load model and processor
6model_name = "Cognitive-Lab/NetraEmbed"
7
8# Load model once (supports all Matryoshka dimensions)
9model = BiGemma3.from_pretrained(
10 model_name,
11 torch_dtype=torch.bfloat16,
12 device_map="cuda",
13)
14processor = BiGemmaProcessor3.from_pretrained(model_name)
15
16# Load your images
17images = [
18 Image.open("document1.jpg"),
19 Image.open("document2.jpg"),
20]
21
22# Define queries
23queries = [
24 "What is the total revenue?",
25 "Show me the organizational chart",
26]
27
28# Process and encode
29batch_images = processor.process_images(images).to(model.device)
30batch_queries = processor.process_texts(queries).to(model.device)
31
32# Choose embedding dimension at inference time: 768, 1536, or 2560
33# Use lower dims for faster search, higher for better accuracy
34embedding_dim = 1536 # Balanced performance
35
36with torch.no_grad():
37 image_embeddings = model(**batch_images, embedding_dim=embedding_dim) # Shape: (num_images, embedding_dim)
38 query_embeddings = model(**batch_queries, embedding_dim=embedding_dim) # Shape: (num_queries, embedding_dim)
39
40# Compute similarity scores using cosine similarity
41scores = processor.score(
42 qs=query_embeddings,
43 ps=image_embeddings,
44) # Shape: (num_queries, num_images)
45
46# Get best matches
47for i, query in enumerate(queries):
48 best_idx = scores[i].argmax().item()
49 print(f"Query: '{query}' -> Best match: Image {best_idx + 1} (score: {scores[i, best_idx]:.4f})")
1# Load model once
2model = BiGemma3.from_pretrained(
3 model_name,
4 torch_dtype=torch.bfloat16,
5 device_map="cuda",
6)
7
8# Test all Matryoshka dimensions
9for embedding_dim in [768, 1536, 2560]:
10 print(f"\nTesting dimension: {embedding_dim}")
11
12 with torch.no_grad():
13 image_embeddings = model(**batch_images, embedding_dim=embedding_dim)
14 query_embeddings = model(**batch_queries, embedding_dim=embedding_dim)
15
16 scores = processor.score(qs=query_embeddings, ps=image_embeddings)
17 print(f"Scores shape: {scores.shape}")
18 print(f"Best match score: {scores.max().item():.4f}")
NetraEmbed achieves state-of-the-art performance on multilingual document retrieval benchmarks. Evaluated on
Nayana-IR Bench (22 languages) and ViDoRe v2.
See our
paper for comprehensive evaluation and per-language analysis.
1@misc{kolavi2025m3druniversalmultilingualmultimodal,
2 title={M3DR: Towards Universal Multilingual Multimodal Document Retrieval},
3 author={Adithya S Kolavi and Vyoman Jain},
4 year={2025},
5 eprint={2512.03514},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2512.03514}
9}
This model is released under the same license as the base Gemma3 model.
This work benefited from compute credits for training, inference, and evaluation provided by
Modal, acknowledged as a compute sponsor. Dataset curation and synthesis were supported by the
Meta LLaMA Impact Grant through our
Nayana initiative. We appreciate Meta for continued support of our research efforts at
CognitiveLab.
Built on top of the Gemma3 architecture with Matryoshka representation learning.
