Update app.py

app.py

@@ -1,27 +1,21 @@
-from transformers import AutoTokenizer, AutoModel
-import torch
 import gradio as gr
+from sentence_transformers import SentenceTransformer
 
-# Load the pre-trained paraphrase-mpnet-base-v2 model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/paraphrase-mpnet-base-v2')
-model = AutoModel.from_pretrained('sentence-transformers/paraphrase-mpnet-base-v2')
+# Load the pre-trained paraphrase-mpnet-base-v2 model
+model = SentenceTransformer('sentence-transformers/paraphrase-mpnet-base-v2')
 
-def get_mpnet_embeddings(sentences):
-    # Tokenize input sentences
-    inputs = tokenizer(sentences, return_tensors='pt', padding=True, truncation=True, max_length=512)
-    # Get embeddings
-    with torch.no_grad():
-        outputs = model(**inputs)
-    embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling over the sequence
-    return embeddings.numpy().tolist()
+def get_embeddings(sentences):
+    # Get embeddings for the input sentences
+    embeddings = model.encode(sentences)
+    return embeddings.tolist()
 
 # Define the Gradio interface
 interface = gr.Interface(
-    fn=get_mpnet_embeddings,  # Function to call
+    fn=get_embeddings,  # Function to call
     inputs=gr.Textbox(lines=2, placeholder="Enter sentences here, one per line"),  # Input component
     outputs=gr.JSON(),  # Output component
     title="Sentence Embeddings with MPNet",  # Interface title
-    description="Enter sentences to get their embeddings with paraphrase-mpnet-base-v2 (up to 512 tokens)."  # Description
+    description="Enter sentences to get their embeddings with paraphrase-mpnet-base-v2."  # Description
 )
 
 # Launch the interface