File size: 260 Bytes
7dff9c2
 
 
 
 
 
 
1
2
3
4
5
6
7
8
from sentence_transformers import SentenceTransformer

def load_model(device="cpu"):
    return SentenceTransformer("all-MiniLM-L6-v2").to(device)

def compute_query_embedding(model, query):
    return model.encode(query, convert_to_tensor=True).cpu().numpy()