Update app.py

app.py

@@ -51,29 +51,29 @@ 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):
+    def compute_embeddings(self, selected_task, input_text):
         max_length = 2042
         task_description = tasks[selected_task]
         processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
     
-        batch_dict = tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
-        batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
-        batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
+        batch_dict = self.tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
+        batch_dict['input_ids'] = [input_ids + [self.tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+        batch_dict = self.tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
         batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
-        outputs = model(**batch_dict)
+        outputs = self.model(**batch_dict)
         embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
         embeddings = F.normalize(embeddings, p=2, dim=1)
         embeddings_list = embeddings.detach().cpu().numpy().tolist()
         return embeddings_list
 
     @spaces.GPU
-    def compute_similarity(self, sentence1, sentence2, extra_sentence1, extra_sentence2):
+    def compute_similarity(self, selected_task, sentence1, sentence2, extra_sentence1, extra_sentence2):
 
         # 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)
+        embeddings1 = self.compute_embeddings(self.selected_task, sentence1)
+        embeddings2 = self.compute_embeddings(self.selected_task, sentence2)
+        embeddings3 = self.compute_embeddings(self.selected_task, extra_sentence1)
+        embeddings4 = self.compute_embeddings(self.selected_task, extra_sentence2)
     
         # Convert embeddings to tensors
         embeddings_tensor1 = torch.tensor(embeddings1).to(device)
@@ -89,6 +89,7 @@ class EmbeddingModel:
 
     
 def app_interface():
+#   embedding_model = EmbeddingModel()
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
@@ -114,7 +115,7 @@ def app_interface():
             similarity_output = gr.Label(label="🐣e5-mistral🛌🏻 Similarity Scores")
             similarity_button.click(
                 fn=EmbeddingModel.compute_similarity,
-                inputs=[sentence1_box, sentence2_box, extra_sentence1_box, extra_sentence2_box],
+                inputs=[task_dropdown, sentence1_box, sentence2_box, extra_sentence1_box, extra_sentence2_box],
                 outputs=similarity_output
             )