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
language = "Sanskrit"
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_idx = inputs["attention_mask"].sum(dim=1) - 1
embedding = outputs.hidden_states[-1][torch.arange(last_token_idx.size(0)), last_token_idx]
embedding = embedding / embedding.norm(dim=-1, keepdim=True)