embed / app.py
codeblacks's picture
Update app.py
c44dddd verified
raw
history blame
842 Bytes
from sentence_transformers import SentenceTransformer
import gradio as gr
import update_packages
import numpy as np
# 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.tolist(), shape # Convert numpy array to list
# Define the Gradio interface
interface = gr.Interface(
fn=generate_embeddings,
inputs=gr.Textbox(lines=5, placeholder="Enter text chunks here...", type="list"),
outputs=[gr.JSON(label="Embeddings"), gr.Label(label="Shape")],
title="Sentence Transformer Embeddings",
description="Generate embeddings for input text chunks."
)
# Launch the Gradio app
interface.launch()