Spaces:
Sleeping
Sleeping
File size: 574 Bytes
e02136d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
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() |