gemma-2-mitra-e

Model Card for gemma2-mitra-embedding

Multilingual sentence embedding model based on Gemma 2, designed for semantic similarity and retrieval. It is used in the Mitra alignment stack to embed sentences in languages such as Sanskrit, Tibetan, Pali, Chinese, and English for cross-lingual sentence alignment and similarity search.

Direct Use

• Semantic similarity: Encode sentences and compare them via cosine similarity (embeddings are L2-normalized).
• Retrieval: Encode queries with the query template and corpus passages with the corpus (no instruction) format; retrieve by nearest-neighbor search (e.g. FAISS).
• Multilingual alignment: Used in this repo to embed source and target sentences for sentence-level alignment (e.g. Buddhist texts across Sanskrit, Tibetan, Pali, Chinese, English).

Downstream Use [optional]

• Integration into alignment pipelines (e.g. Bertalign-style alignment with this embedder).
• RAG or search systems that need multilingual, instruction-aware query/corpus embeddings.
• Any application that consumes L2-normalized sentence vectors from this model.

Recommendations

• Use the exact prompt format (see “How to Get Started”) for queries and corpus.
• Users should be aware of potential biases and limitations; evaluate on their own data and languages before deployment.

How to Get Started with the Model

Prompt / template pattern

The model expects asymmetric inputs:

• Query (instruction + query):
  <instruct>Please find the semantically most similar text in {language}.\n<query>{sentence_text}
  where {language} is a full language name (e.g. "Sanskrit", "Tibetan", "English", "Chinese", "Pali", "Hindi") and {sentence_text} is the query string.

• Corpus:
  Use the raw sentence (or passage) text only, with no <instruct> or <query> wrapper.

Example (with 8-bit and Hugging Face Transformers)

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

model_path = "gemma2-mitra-embedding"
quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
    llm_int8_has_fp16_weight=False,
)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    quantization_config=quantization_config,
    device_map={"": 0},
    torch_dtype=torch.float16,
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Query encoding (one sentence)
language = "Sanskrit"  # or Tibetan, English, Chinese, Pali, Hindi
text = "Your query sentence here."
prompt = f"<instruct>Please find the semantically most similar text in {language}.\\n<query>{text}"
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
with torch.no_grad():
    outputs = model(**inputs, output_hidden_states=True)
# Last token embedding (last non-padded)
last_token_idx = inputs["attention_mask"].sum(dim=1) - 1
embedding = outputs.hidden_states[-1][torch.arange(last_token_idx.size(0)), last_token_idx]
# L2-normalize for cosine similarity
embedding = embedding / embedding.norm(dim=-1, keepdim=True)

For corpus sentences, pass only the raw text (no <instruct>/<query>), then take the last token hidden state and L2-normalize the same way.

Alternatively, use FlagEmbedding’s FlagLLMModel with this model path for encode_queries and encode_corpus (see FlagEmbedding).

Model Architecture and Objective

• Architecture: Gemma 2 (Gemma2Model), used as an encoder: input text → last-token hidden state → L2-normalized embedding.
• Config (from repo): hidden_size=3584, num_hidden_layers=42, num_attention_heads=16, num_key_value_heads=8, intermediate_size=14336, head_dim=256, max_position_embeddings=8192, sliding_window=4096, vocab_size=256002 (includes special tokens <instruct>, <query>).
• Special tokens: <instruct>, <query> (see special_tokens_map.json / added_tokens.json in the model dir).
• Objective: Dense retrieval / semantic similarity (asymmetric query/corpus encoding).