Views
No views yet
<instruct>Please find the semantically most similar text in {language}.\n<query>{sentence_text}{language} is a full language name (e.g. "Sanskrit", "Tibetan", "English", "Chinese", "Pali", "Hindi") and {sentence_text} is the query string.<instruct> or <query> wrapper.1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2import torch
3
4model_path = "gemma2-mitra-embedding"
5quantization_config = BitsAndBytesConfig(
6 load_in_8bit=True,
7 llm_int8_threshold=6.0,
8 llm_int8_has_fp16_weight=False,
9)
10model = AutoModelForCausalLM.from_pretrained(
11 model_path,
12 quantization_config=quantization_config,
13 device_map={"": 0},
14 torch_dtype=torch.float16,
15 trust_remote_code=True,
16)
17tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
18if tokenizer.pad_token is None:
19 tokenizer.pad_token = tokenizer.eos_token
20
21# Query encoding (one sentence)
22language = "Sanskrit" # or Tibetan, English, Chinese, Pali, Hindi
23text = "Your query sentence here."
24prompt = f"<instruct>Please find the semantically most similar text in {language}.\\n<query>{text}"
25inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
26with torch.no_grad():
27 outputs = model(**inputs, output_hidden_states=True)
28# Last token embedding (last non-padded)
29last_token_idx = inputs["attention_mask"].sum(dim=1) - 1
30embedding = outputs.hidden_states[-1][torch.arange(last_token_idx.size(0)), last_token_idx]
31# L2-normalize for cosine similarity
32embedding = embedding / embedding.norm(dim=-1, keepdim=True)<instruct>/<query>), then take the last token hidden state and L2-normalize the same way.FlagLLMModel with this model path for encode_queries and encode_corpus (see FlagEmbedding).Gemma2Model), used as an encoder: input text → last-token hidden state → L2-normalized embedding.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>).<instruct>, <query> (see special_tokens_map.json / added_tokens.json in the model dir).