Update app.py

app.py

@@ -51,7 +51,7 @@ class EmbeddingModel:
         self.model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype=torch.float16, device_map=device)
 
     @spaces.GPU
-    def compute_embeddings(selected_task, input_text, system_prompt):
+    def compute_embeddings(selected_task, input_text):
         max_length = 2042
         task_description = tasks[selected_task]
         processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
@@ -69,17 +69,22 @@ class EmbeddingModel:
     @spaces.GPU
     def compute_similarity(self, sentence1, sentence2, extra_sentence1, extra_sentence2):
 
-        sentences = [sentence1, sentence2, extra_sentence1, extra_sentence2]
-        encoded_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt').to(device)
-        with torch.no_grad():
-            model_output = self.model(**encoded_input)
-
-        embeddings = last_token_pool(model_output.last_hidden_state, encoded_input['attention_mask'])
-        embeddings = F.normalize(embeddings, p=2, dim=1)
-
+        # Compute embeddings for each sentence
+        embeddings1 = compute_embeddings(self.selected_task, sentence1)
+        embeddings2 = compute_embeddings(self.selected_task, sentence2)
+        embeddings3 = compute_embeddings(self.selected_task, extra_sentence1)
+        embeddings4 = compute_embeddings(self.selected_task, extra_sentence2)
+    
+        # Convert embeddings to tensors
+        embeddings_tensor1 = torch.tensor(embeddings1).to(device)
+        embeddings_tensor2 = torch.tensor(embeddings2).to(device)
+        embeddings_tensor3 = torch.tensor(embeddings3).to(device)
+        embeddings_tensor4 = torch.tensor(embeddings4).to(device)
+    
         # Compute cosine similarity
-        similarity1 = F.cosine_similarity(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0)).item()
-        similarity2 = F.cosine_similarity(embeddings[2].unsqueeze(0), embeddings[3].unsqueeze(0)).item()
+        similarity1 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor2).item()
+        similarity2 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor3).item()
+        similarity3 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor4).item()
         return similarity1, similarity2
 
     
@@ -92,20 +97,19 @@ def app_interface():
 
         with gr.Tab("Embedding Generation"):
             input_text_box = gr.Textbox(label="📖Input Text")
-            system_prompt_box = gr.Textbox(label="🤖System Prompt (Optional)")
             compute_button = gr.Button("Try🐣🛌🏻e5")
             output_display = gr.Textbox(label="🐣e5-mistral🛌🏻 Embeddings")
             compute_button.click(
                 fn=EmbeddingModel.compute_embeddings,
-                inputs=[task_dropdown, input_text_box, system_prompt_box],
+                inputs=[task_dropdown, input_text_box],
                 outputs=output_display
             )
 
         with gr.Tab("Sentence Similarity"):
-            sentence1_box = gr.Textbox(label="Sentence 1")
-            sentence2_box = gr.Textbox(label="Sentence 2")
-            extra_sentence1_box = gr.Textbox(label="Sentence 3")
-            extra_sentence2_box = gr.Textbox(label="Sentence 4")
+            sentence1_box = gr.Textbox(label="'Focus Sentence' - The 'Subject'")
+            sentence2_box = gr.Textbox(label="'Input Sentence' - 1")
+            extra_sentence1_box = gr.Textbox(label="'Input Sentence' - 2")
+            extra_sentence2_box = gr.Textbox(label="'Input Sentence' - 3")
             similarity_button = gr.Button("Compute Similarity")
             similarity_output = gr.Label(label="🐣e5-mistral🛌🏻 Similarity Scores")
             similarity_button.click(
@@ -116,7 +120,6 @@ def app_interface():
 
         with gr.Row():
             with gr.Column():
-                system_prompt_box
                 input_text_box
             with gr.Column():
                 compute_button