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()