TechKick's picture
Make embeddings compatible with ZeroGPU
eb850ac
Raw
History Blame Contribute Delete
250 Bytes
import spaces
from sentence_transformers import SentenceTransformer
model = None
@spaces.GPU
def get_embeddings(texts):
global model
if model is None:
model = SentenceTransformer("all-MiniLM-L6-v2")
return model.encode(texts)