from sentence_transformers import SentenceTransformer
import gradio as gr
import update_packages
import numpy as np

# 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=False)
    shape = embeddings.shape
    return embeddings.tolist(), shape  # Convert numpy array to list

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

# Launch the Gradio app
interface.launch()