Update app.py

app.py

@@ -1,14 +1,14 @@
-from transformers import LongformerTokenizer, LongformerModel
+from transformers import AutoTokenizer, AutoModel
 import torch
 import gradio as gr
 
-# Load the pre-trained Longformer model and tokenizer
-tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
-model = LongformerModel.from_pretrained('allenai/longformer-base-4096')
+# 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')
 
-def get_longformer_embeddings(sentences):
+def get_mpnet_embeddings(sentences):
     # Tokenize input sentences
-    inputs = tokenizer(sentences, return_tensors='pt', padding=True, truncation=True, max_length=2048)
+    inputs = tokenizer(sentences, return_tensors='pt', padding=True, truncation=True, max_length=512)
     # Get embeddings
     with torch.no_grad():
         outputs = model(**inputs)
@@ -17,11 +17,11 @@ def get_longformer_embeddings(sentences):
 
 # Define the Gradio interface
 interface = gr.Interface(
-    fn=get_longformer_embeddings,  # Function to call
+    fn=get_mpnet_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 Longformer",  # Interface title
-    description="Enter sentences to get their embeddings with Longformer (up to 2048 tokens)."  # Description
+    title="Sentence Embeddings with MPNet",  # Interface title
+    description="Enter sentences to get their embeddings with paraphrase-mpnet-base-v2 (up to 512 tokens)."  # Description
 )
 
 # Launch the interface