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=False)
    shape= embeddings.shape
    return embeddings, shape  # Convert tensor to list for Gradio

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

# Launch the Gradio app
interface.launch()