Spaces:
Sleeping
Sleeping
from transformers import AutoTokenizer, AutoModel | |
import torch | |
def get_embedder(): | |
model_name = "microsoft/MiniLM-L12-H384-uncased" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModel.from_pretrained(model_name) | |
return tokenizer, model | |
def embed_text(texts, tokenizer, model): | |
encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors="pt") | |
with torch.no_grad(): | |
model_output = model(**encoded_input) | |
embeddings = model_output.last_hidden_state.mean(dim=1) | |
return embeddings.numpy().tolist() |