Update app.py

app.py

@@ -1,36 +1,22 @@
 from sentence_transformers import SentenceTransformer
 import gradio as gr
 import torch
-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(text):
-    # Split the input text into chunks (if needed)
-    chunks = text.split('\n')  # Assuming chunks are separated by new lines
-
-    # Encode the input chunks to get the embeddings
-    embeddings = embedding_model.encode(chunks, convert_to_tensor=False)
-
-    # Convert the embeddings to a PyTorch tensor
-    embeddings_tensor = torch.tensor(embeddings)
-
-    # Add batch dimension to the tensor (if needed)
-    embeddings_tensor = embeddings_tensor.unsqueeze(0)  # Uncomment if a batch dimension is required
-
-    # Return the embeddings tensor and its shape
-    return embeddings_tensor.tolist(), embeddings_tensor.shape
+def get_embeddings(sentences):
+    embeddings = model.encode(sentences, convert_to_tensor=True)
+    return embeddings.tolist()
 
 # Define the Gradio interface
 interface = gr.Interface(
-    fn=generate_embeddings,
-    inputs=gr.Textbox(lines=5, placeholder="Enter text chunks here..."),
-    outputs=[gr.JSON(label="Embeddings"), gr.Label(label="Shape")],
-    title="Sentence Transformer Embeddings",
-    description="Generate embeddings for input text chunks."
+    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 Gradio app
-interface.launch()
+# Launch the interface
+interface.launch()