Update app.py

app.py

@@ -1,21 +1,32 @@
 from sentence_transformers import SentenceTransformer
 import gradio as gr
-import update_packages
+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(chunks):
+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)
-    shape = embeddings.shape
-    return embeddings.tolist(), shape  # Convert numpy array to list
+
+    # 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
 
 # Define the Gradio interface
 interface = gr.Interface(
     fn=generate_embeddings,
-    inputs=gr.Textbox(lines=5, placeholder="Enter text chunks here...", type="text"),
+    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."