Update app.py

app.py

@@ -52,8 +52,12 @@ class EmbeddingModel:
 
     @spaces.GPU
     def compute_embeddings(self, selected_task, input_text):
+        try:
+            task_description = tasks[selected_task]
+        except KeyError:
+            print(f"Selected task not found: {selected_task}")
+            return f"Error: Task '{selected_task}' not found. Please select a valid task."
         max_length = 2042
-        task_description = tasks[selected_task]
         processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
     
         batch_dict = self.tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
@@ -68,7 +72,11 @@ class EmbeddingModel:
 
     @spaces.GPU
     def compute_similarity(self, selected_task, sentence1, sentence2, extra_sentence1, extra_sentence2):
-
+        try:
+            task_description = tasks[selected_task]
+        except KeyError:
+            print(f"Selected task not found: {selected_task}")
+            return f"Error: Task '{selected_task}' not found. Please select a valid task."
         # Compute embeddings for each sentence
         embeddings1 = self.compute_embeddings(self.selected_task, sentence1)
         embeddings2 = self.compute_embeddings(self.selected_task, sentence2)