Update app.py

app.py

@@ -8,6 +8,11 @@ from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel, Field
 from rag_util import Create_RAG_Prompt, ProcessDocuments
 import time
+import torch
+
+
+torch.set_num_threads(os.cpu_count())
+
 
 class DataModel(BaseModel):
     data: str
@@ -44,7 +49,7 @@ async def generate(data: DataModel):
         # input_text = bleach.clean(data.data)
         history = chat_history[data.chatID] = [] if data.chatID not in chat_history else chat_history[data.chatID]
 
-        llm_response = generate_response(prompt=input_text, history=history)
+        llm_response = await generate_response(prompt=input_text, history=history)
         
         if data.chatID not in chat_history:
             chat_history[data.chatID] = []
@@ -57,7 +62,7 @@ async def generate(data: DataModel):
         # input_text = bleach.clean(data.data)
         prompt, context = Create_RAG_Prompt(input_text, chatID=data.chatID, history=chat_history[data.chatID] if chat_history[data.chatID] else [])
 
-        llm_response = generate_response(prompt=prompt, context=context)
+        llm_response = await generate_response(prompt=prompt, context=context)
        
         if data.chatID not in chat_history:
             chat_history[data.chatID] = []