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

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

def get_embeddings(sentences):
    # Split sentences by new line
    # sentences_list = [s.strip() for s in sentences.split('\n') if s.strip()]
    # Get embeddings for the input sentences
    embeddings = model.encode(sentences, convert_to_tensor=True)
    # Convert to 2D NumPy array
    # embeddings_array = np.array(embeddings)
    embeddings_array=embeddings.tolist()
    return embeddings_array

# Define the Gradio interface
interface = gr.Interface(
    fn=get_embeddings,  # Function to call
    inputs=gr.Textbox(lines=2, placeholder="Enter sentences here, one per line"),  # Input component
    outputs=gr.Array(),  # Output component
    title="Sentence Embeddings",  # Interface title
    description="Enter sentences to get their embeddings."  # Description
)

# Launch the interface
interface.launch()