Update app.py

app.py

@@ -3,6 +3,8 @@ import torch
 import torch.nn.functional as F
 from torch import Tensor
 from transformers import AutoTokenizer, AutoModel
+import threading
+import queue
 import gradio as gr
 import os
 
@@ -33,6 +35,12 @@ tasks = {
         'TRECCOVID': 'Given a query on COVID-19, retrieve documents that answer the query',
 }
 
+
+# Global queue for embedding requests
+embedding_request_queue = queue.Queue()
+embedding_response_queue = queue.Queue()
+
+
 tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-mistral-7b-instruct')
 model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype=torch.float16, device_map=device)
 
@@ -56,7 +64,24 @@ def load_corpus_from_json(file_path):
     with open(file_path, 'r') as file:
         data = json.load(file)
     return data
-    
+
+
+def embedding_worker():
+    while True:
+        # Wait for an item in the queue
+        item = embedding_request_queue.get()
+        if item is None:
+            break
+        selected_task, input_text = item
+        embeddings = compute_embeddings(selected_task, input_text)
+        formatted_response = format_response(embeddings)
+
+        embedding_response_queue.put(formatted_response)
+        embedding_request_queue.task_done()
+
+threading.Thread(target=embedding_worker, daemon=True).start()
+
+
 @spaces.GPU
 def compute_embeddings(selected_task, input_text):
     try:
@@ -160,8 +185,11 @@ def format_response(embeddings):
     }
 
 def generate_and_format_embeddings(selected_task, input_text):
-    embeddings = compute_embeddings(selected_task, input_text)
-    return format_response(embeddings)
+    embedding_request_queue.put((selected_task, input_text))
+    response = embedding_response_queue.get()
+    embedding_response_queue.task_done()
+    return response
+
 
 def app_interface():
     corpus_sentences = []