# app.py
import gradio as gr
from sentence_transformers import SentenceTransformer

# Initialize the embedding model
model = SentenceTransformer('distilbert-base-nli-stsb-mean-tokens')

def get_embedding(text):
    embedding = model.encode(text)
    return embedding.tolist()

# Define the Gradio interface using updated components
iface = gr.Interface(
    fn=get_embedding,
    inputs=gr.Textbox(lines=5, placeholder="Enter text here..."),
    outputs=gr.JSON(),
    title="Embedding Generation Service",
    description="Generates sentence embeddings using DistilBERT."
)

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