Spaces:
Sleeping
Sleeping
| 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() | |