Update app.py

app.py

@@ -1,14 +1,13 @@
 import streamlit as st
 import hashlib
 import os
-import requests
+import aiohttp
+import asyncio
 import time
 from langsmith import traceable
 import random
 import discord
-import os
 from transformers import pipeline
-import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
@@ -18,8 +17,6 @@ from tqdm import tqdm
 import re
 import os
 
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
 st.set_page_config(page_title="TeapotAI Discord Bot", page_icon=":robot_face:", layout="wide")
 tokenizer = None
 model = None
@@ -28,11 +25,10 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
 
-
 def log_time(func):
-    def wrapper(*args, **kwargs):
+    async def wrapper(*args, **kwargs):
         start_time = time.time()
-        result = func(*args, **kwargs)
+        result = await func(*args, **kwargs)  # Make it awaitable
         end_time = time.time()
         print(f"{func.__name__} executed in {end_time - start_time:.4f} seconds")
         return result
@@ -42,24 +38,25 @@ def log_time(func):
 API_KEY = os.environ.get("brave_api_key")
 
 @log_time
-def brave_search(query, count=3):
+async def brave_search(query, count=3):
     url = "https://api.search.brave.com/res/v1/web/search"
     headers = {"Accept": "application/json", "X-Subscription-Token": API_KEY}
     params = {"q": query, "count": count}
-    
-    response = requests.get(url, headers=headers, params=params)
-    
-    if response.status_code == 200:
-        results = response.json().get("web", {}).get("results", [])
-        print(results)
-        return [(res["title"], res["description"], res["url"]) for res in results]
-    else:
-        print(f"Error: {response.status_code}, {response.text}")
-        return []
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url, headers=headers, params=params) as response:
+            if response.status == 200:
+                results = await response.json()
+                print(results)
+                return [(res["title"], res["description"], res["url"]) for res in results.get("web", {}).get("results", [])]
+            else:
+                print(f"Error: {response.status}, {await response.text()}")
+                return []
+
 
 @traceable 
 @log_time
-def query_teapot(prompt, context, user_input):
+async def query_teapot(prompt, context, user_input):
     input_text = prompt + "\n" + context + "\n" + user_input
     print(input_text)
     start_time = time.time()
@@ -67,7 +64,7 @@ def query_teapot(prompt, context, user_input):
     inputs = tokenizer(input_text, return_tensors="pt")
     input_length = inputs["input_ids"].shape[1]
     
-    output = model.generate(**inputs, max_new_tokens=512)
+    output = await asyncio.to_thread(model.generate, **inputs, max_new_tokens=512)
     
     output_text = tokenizer.decode(output[0], skip_special_tokens=True)
     total_length = output.shape[1]  # Includes both input and output tokens
@@ -81,19 +78,18 @@ def query_teapot(prompt, context, user_input):
     return output_text
 
 
-
 @log_time
-def handle_chat(user_input):
+async def handle_chat(user_input):
     search_start_time = time.time()
-    results = brave_search(user_input)
+    results = await brave_search(user_input)
     search_end_time = time.time()
     
-    documents = [desc.replace('<strong>','').replace('</strong>','') for _, desc, _ in results]
+    documents = [desc.replace('<strong>', '').replace('</strong>', '') for _, desc, _ in results]
 
     context = "\n".join(documents)
     prompt = """You are Teapot, an open-source AI assistant optimized for low-end devices, providing short, accurate responses without hallucinating while excelling at information extraction and text summarization. If a user asks who you are reply "I am Teapot"."""
     generation_start_time = time.time()
-    response = query_teapot(prompt, context, user_input)
+    response = await query_teapot(prompt, context, user_input)
     generation_end_time = time.time()
     
     debug_info = f"""
@@ -108,9 +104,9 @@ Generation time: {generation_end_time - generation_start_time:.2f} seconds
 Response: {response}
 """
     
-    
     return response, debug_info
 
+
 st.write("418 I'm a teapot")
 
 DISCORD_TOKEN = os.environ.get("discord_key")
@@ -135,11 +131,10 @@ async def on_message(message):
         return
     print(message.content)
 
-
     is_debug = "<debug>" in message.content
     async with message.channel.typing():
         # Respond with "pong" if the message contains "ping"
-        response, debug_info = handle_chat(message.content.replace("<debug>","").replace("</debug>",""))
+        response, debug_info = await handle_chat(message.content.replace("<debug>", "").replace("</debug>", ""))
         
         print(response)
         sent_message = await message.reply(response)
@@ -153,6 +148,4 @@ async def on_message(message):
 
 
 # Run the bot with your token
-
 client.run(DISCORD_TOKEN)
-