Spaces:
Sleeping
Sleeping
File size: 522 Bytes
a67d009 226e778 a67d009 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np
MODEL = "sdadas/stella-pl-retrieval-8k"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModel.from_pretrained(MODEL)
def embed(text):
tokens = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
out = model(**tokens).last_hidden_state
emb = out.mean(dim=1).squeeze().tolist()
return emb
gr.Interface(embed, gr.Textbox(), gr.JSON()).launch()
|