Update app.py

app.py

@@ -122,17 +122,14 @@ def compute_embeddings(selected_task, input_text):
         print(f"Selected task not found: {selected_task}")
         return f"Error: Task '{selected_task}' not found. Please select a valid task."
     
-    max_length = 2048
-    processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
+    query_prefix = f"Instruct: {task_description}\nQuery: "
+    queries = [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 = {k: v.to(device) for k, v in batch_dict.items()}
-    outputs = 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()
+    # Get the embeddings
+    query_embeddings = model.encode(queries, instruction=query_prefix, max_length=4096)
+    # Normalize embeddings
+    query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
+    embeddings_list = query_embeddings.detach().cpu().numpy().tolist()
     clear_cuda_cache()
     return embeddings_list