embed / app.py
codeblacks's picture
Update app.py
18feec4 verified
raw
history blame
748 Bytes
from sentence_transformers import SentenceTransformer
import gradio as gr
import torch
# Load the pre-trained model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
def get_embeddings(sentences):
embeddings = model.encode(sentences, convert_to_tensor=True)
return embeddings.tolist()
# 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.Image(label="Embeddings", image_formatter=plot_embeddings)
title="Sentence Embeddings", # Interface title
description="Enter sentences to get their embeddings." # Description
)
# Launch the interface
interface.launch()