Add asynchronous responses

app.py

@@ -8,6 +8,8 @@ import gspread
 from oauth2client.service_account import ServiceAccountCredentials
 import gradio as gr
 from huggingface_hub import login
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 
 
 BOOK_MAPPING = {
@@ -39,26 +41,25 @@ class ModelManager:
         self.model_B = Client("mep296/llama-3-8b-rephrase-summarize-quality")
         self.model_C = Client("mep296/llama-3-8b-entigraph-quality")
         self.template_text = self._load_template()
-    
+
     def _load_template(self):
         with open("prompt_template.txt", "r", encoding="utf-8") as file:
             return file.read()
-    
-    def get_model_response(self, model_name, prompt):
+
+    async def get_model_response_async(self, model_name, client, prompt):
         try:
             formatted_prompt = self.template_text.format(prompt)
             
-            model_clients = {
-                "Model_A": self.model_A,
-                "Model_B": self.model_B,
-                "Model_C": self.model_C
-            }
+            loop = asyncio.get_running_loop()
+            with ThreadPoolExecutor() as executor:
+                response = await loop.run_in_executor(
+                    executor, 
+                    client.predict, 
+                    formatted_prompt, 
+                    "/predict"
+                )
             
-            client = model_clients[model_name]
-            response = client.predict(
-                prompt=formatted_prompt,
-                api_name="/predict"
-            )
+            model_responses[model_name] = response
             return response
         except (httpx.ReadTimeout, httpx.ConnectTimeout) as e:
             print(f"Timeout while getting response from {model_name}: {str(e)}")
@@ -66,19 +67,15 @@ class ModelManager:
         except Exception as e:
             print(f"Error getting response from {model_name}: {str(e)}")
             return f"Error: Could not get response from {model_name}. Please try again."
-    
-    def get_all_model_responses(self, prompt):
-        responses = []
-        model_responses.clear()
 
-        self.model_A = Client("mep296/llama-3-8b-rephrase-quality")
-        self.model_B = Client("mep296/llama-3-8b-rephrase-summarize-quality")
-        self.model_C = Client("mep296/llama-3-8b-entigraph-quality")
-
-        for model in ['A', 'B', 'C']:
-            response = self.get_model_response(f"Model_{model}", prompt)
-            model_responses[f"Model_{model}"] = response
-            responses.append(response)
+    async def get_all_model_responses_async(self, prompt):
+        tasks = [
+            self.get_model_response_async("Model_A", self.model_A, prompt),
+            self.get_model_response_async("Model_B", self.model_B, prompt),
+            self.get_model_response_async("Model_C", self.model_C, prompt)
+        ]
+        
+        responses = await asyncio.gather(*tasks)
         return responses
 
     def get_book_model_mapping(self, book):
@@ -128,6 +125,7 @@ class ModelComparisonApp:
         self.chat_history_B = []
         self.chat_history_C = []
         self.state = gr.State(value="")
+        self.loop = asyncio.get_event_loop()
         
     def create_interface(self):
         text_size = gr.themes.sizes.text_lg
@@ -191,7 +189,10 @@ class ModelComparisonApp:
         mapping_dict = self.model_manager.book_model_assignments[self.selected_book]
         model_order = ["rephrase", "rephrase_summarize", "entigraph"]
         model_to_index = {model: i for i, model in enumerate(model_order)}
-        responses = self.model_manager.get_all_model_responses(prompt)
+
+        responses = self.loop.run_until_complete(
+            self.model_manager.get_all_model_responses_async(prompt)
+        )
         
         chats = []
         for response in responses:
@@ -202,7 +203,7 @@ class ModelComparisonApp:
         
         reordered_chats = [chats[model_to_index[mapping_dict[model]]] for model in ["Model A", "Model B", "Model C"]]
         return reordered_chats
-    
+
     def get_votes(self, book, category, question, rating_A, rating_B, rating_C):
         model_mapping = self.model_manager.get_book_model_mapping(book)
         rating_data = {