import gradio as gr
from sentence_transformers import SentenceTransformer

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

def get_embeddings(sentences):
    # Check if the input is a list
    if isinstance(sentences, list):
        sentence_list = sentences
    elif isinstance(sentences, str):
        # If it's a string, split by new lines to create a list of sentences
        sentence_list = sentences.split("\n")
    else:
        raise ValueError("Input should be either a string or a list of strings.")
    
    # Generate embeddings
    embeddings = model.encode(sentence_list, convert_to_tensor=True)
    return str(embeddings.tolist())

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

# Launch the interface
interface.launch()