import json
import gradio as gr
import spaces
from sentence_transformers import SentenceTransformer

model = None

def get_model():
    global model
    if model is None:
        model = SentenceTransformer("mixedbread-ai/deepset-mxbai-embed-de-large-v1")
    return model

@spaces.GPU
def embed_texts(texts_json: str) -> str:
    """Embed a list of texts. Input: JSON array of strings. Output: JSON array of float arrays."""
    texts = json.loads(texts_json)
    m = get_model()
    embeddings = m.encode(texts, batch_size=64, convert_to_numpy=True)
    return json.dumps(embeddings.tolist())

demo = gr.Interface(
    fn=embed_texts,
    inputs=gr.Textbox(label="Texts (JSON array of strings)", lines=5),
    outputs=gr.Textbox(label="Embeddings (JSON array of float arrays)", lines=5),
    title="HDP Wiki Embedder",
    description="Batch text embedding using deepset-mxbai-embed-de-large-v1 (1024-dim).",
)

if __name__ == "__main__":
    demo.launch()