Update app.py

app.py

@@ -1,5 +1,101 @@
+import streamlit as st
+import hashlib
+import os
+import requests
+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
+from pydantic import BaseModel
+from typing import List, Optional
+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
+model_name = "teapotai/teapotllm"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+
+
+def log_time(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        print(f"{func.__name__} executed in {end_time - start_time:.4f} seconds")
+        return result
+    return wrapper
+
+
+API_KEY = os.environ.get("brave_api_key")
+
+@log_time
+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 []
+
+@traceable 
+@log_time
+def query_teapot(prompt, context, user_input):
+    input_text = prompt + "\n" + context + "\n" + user_input
+    
+    start_time = time.time()
+    
+    inputs = tokenizer(input_text, return_tensors="pt")
+    input_length = inputs["input_ids"].shape[1]
+    
+    output = 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
+    output_length = total_length - input_length  # Extract output token count
+    
+    end_time = time.time()
+    
+    elapsed_time = end_time - start_time
+    tokens_per_second = total_length / elapsed_time if elapsed_time > 0 else float("inf")
+    
+    return output_text
+
+
+
+@log_time
+def handle_chat(user_prompt, user_input):
+
+    results = brave_search(user_input)
+    
+    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"."""
+    response = query_teapot(prompt, context+user_prompt, user_input)
+
+    return response
+
+st.write("418 I'm a teapot")
 
 DISCORD_TOKEN = os.environ.get("discord_key")
 
@@ -23,8 +119,8 @@ async def on_message(message):
         return
 
     # Respond with "pong" if the message contains "ping"
-    if 'ping' in message.content.lower():
-        await message.channel.send('pong')
+    response = handle_chat(message.content)
+    await message.channel.send(response)
 
 # Run the bot with your token