Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2import torch
3import torch.nn.functional as F
4
5model = AutoModel.from_pretrained('rbhatia46/harrier-270m-lora-attn-downproj-qwen3-oracle', torch_dtype=torch.bfloat16)
6tok = AutoTokenizer.from_pretrained('rbhatia46/harrier-270m-lora-attn-downproj-qwen3-oracle')
7
8QUERY_INST = 'Given a web search query, retrieve relevant passages that answer the query'
9
10def encode(texts, is_query=False):
11 if is_query:
12 texts = [f'Instruct: {QUERY_INST}\nQuery: {t}' for t in texts]
13 enc = tok(texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
14 with torch.no_grad():
15 out = model(**enc)
16 # Last token pooling
17 seq_lens = enc['attention_mask'].sum(dim=1) - 1
18 emb = out.last_hidden_state[torch.arange(len(seq_lens)), seq_lens]
19 return F.normalize(emb.float(), p=2, dim=-1)