from sentence_transformers import SentenceTransformer
import gradio as gr

# Load the pre-trained model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

# Define the function to process requests
def generate_embeddings(chunks):
    embeddings = embedding_model.encode(chunks, convert_to_tensor=True)
    return embeddings.tolist()  # Convert tensor to list for Gradio

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_embeddings,
    inputs=gr.inputs.Textbox(lines=5, placeholder="Enter text chunks here..."),
    outputs=gr.outputs.JSON(),
    title="Sentence Transformer Embeddings",
    description="Generate embeddings for input text chunks."
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()