codeblacks's picture
Update app.py
a95d821 verified
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()