Update app.py

app.py

@@ -1,155 +1,98 @@
+import os
+import urllib
+import requests
+import feedparser
+from bs4 import BeautifulSoup
+import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import logging
-import random
-import requests
-import urllib
-from bs4 import BeautifulSoup
-import os
 
-# Initialize logging
+# Set up logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# List of user agents to choose from for requests
-_useragent_list = [
-    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:66.0) Gecko/20100101 Firefox/66.0',
-    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
-    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
-    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36',
-    'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
-    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36 Edg/111.0.1661.62',
-    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0'
-]
-
-def get_useragent():
-    """Returns a random user agent from the list."""
-    return random.choice(_useragent_list)
-
-def extract_text_from_webpage(html_content):
-    """Extracts visible text from HTML content using BeautifulSoup."""
-    soup = BeautifulSoup(html_content, "html.parser")
-    # Remove unwanted tags
-    for tag in soup(["script", "style", "header", "footer", "nav"]):
-        tag.extract()
-    # Get the remaining visible text
-    visible_text = soup.get_text(strip=True)
-    return visible_text
-
-def search(term, num_results=1, lang="en", advanced=True, sleep_interval=0, timeout=5, safe="active", ssl_verify=None):
-    """Performs a Google search and returns the results."""
-    escaped_term = urllib.parse.quote_plus(term)
-    start = 0
-    all_results = []
-
-    # Fetch results in batches
-    while start < num_results:
-        resp = requests.get(
-            url="https://www.google.com/search",
-            headers={"User-Agent": get_useragent()}, # Set random user agent
-            params={
-                "q": term,
-                "num": num_results - start, # Number of results to fetch in this batch
-                "hl": lang,
-                "start": start,
-                "safe": safe,
-            },
-            timeout=timeout,
-            verify=ssl_verify,
-        )
-        resp.raise_for_status() # Raise an exception if request fails
-
-        soup = BeautifulSoup(resp.text, "html.parser")
-        result_block = soup.find_all("div", attrs={"class": "g"})
-
-        # If no results, continue to the next batch
-        if not result_block:
-            start += 1
-            continue
-
-        # Extract link and text from each result
-        for result in result_block:
-            link = result.find("a", href=True)
-            if link:
-                link = link["href"]
-                try:
-                    # Fetch webpage content
-                    webpage = requests.get(link, headers={"User-Agent": get_useragent()})
-                    webpage.raise_for_status()
-                    # Extract visible text from webpage
-                    visible_text = extract_text_from_webpage(webpage.text)
-                    all_results.append({"link": link, "text": visible_text})
-                except requests.exceptions.RequestException as e:
-                    # Handle errors fetching or processing webpage
-                    print(f"Error fetching or processing {link}: {e}")
-                    all_results.append({"link": link, "text": None})
-            else:
-                all_results.append({"link": None, "text": None})
-
-        start += len(result_block) # Update starting index for next batch
-
-    return all_results
-
-# Load the model and tokenizer
-model_name = "mistralai/Mistral-7B-v0.3"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to(DEVICE)
-
-def format_prompt(user_input, chat_history):
+# Define device and load model and tokenizer
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+
+# Load model and tokenizer
+try:
+    logger.debug("Attempting to load the model and tokenizer")
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    logger.debug("Model and tokenizer loaded successfully")
+except Exception as e:
+    logger.error(f"Error loading model and tokenizer: {e}")
+    model = None
+    tokenizer = None
+
+# Function to fetch news from Google News RSS feed
+def fetch_news(term, num_results=2):
+    logger.debug(f"Fetching news for term: {term}")
+    url = f"https://news.google.com/rss/search?q={term}"
+    feed = feedparser.parse(url)
+    results = []
+    for entry in feed.entries[:num_results]:
+        results.append({"link": entry.link, "text": entry.title})
+    logger.debug(f"Fetched news results: {results}")
+    return results
+
+# Function to format the prompt for the language model
+def format_prompt(user_prompt, chat_history):
+    logger.debug(f"Formatting prompt with user prompt: {user_prompt} and chat history: {chat_history}")
     prompt = ""
-    for user, bot in chat_history:
-        prompt += f"User: {user}\nBot: {bot}\n"
-    prompt += f"User: {user_input}\nBot: "
+    for item in chat_history:
+        prompt += f"User: {item[0]}\nAssistant: {item[1]}\n"
+    prompt += f"User: {user_prompt}\nAssistant:"
+    logger.debug(f"Formatted prompt: {prompt}")
     return prompt
 
-def model_inference(user_prompt, chat_history, web_search, temperature, max_new_tokens, repetition_penalty, top_p):
-    try:
-        if not user_prompt["files"]:
-            if web_search:
-                logger.debug("Performing web search")
-                web_results = search(user_prompt["text"], num_results=3)  # Fetching more results for better context
-                web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results if res['text']])
-                formatted_prompt = format_prompt(f"{user_prompt['text']} [WEB] {web2}", chat_history)
-                inputs = tokenizer(formatted_prompt, return_tensors="pt").to(DEVICE)
-                if model:
-                    outputs = model.generate(
-                        **inputs,
-                        max_new_tokens=max_new_tokens,
-                        repetition_penalty=repetition_penalty,
-                        do_sample=True,
-                        temperature=temperature,
-                        top_p=top_p
-                    )
-                    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                else:
-                    response = "Model is not available. Please try again later."
-                logger.debug(f"Model response: {response}")
-                return response
-            else:
-                formatted_prompt = format_prompt(user_prompt["text"], chat_history)
-                inputs = tokenizer(formatted_prompt, return_tensors="pt").to(DEVICE)
-                if model:
-                    outputs = model.generate(
-                        **inputs,
-                        max_new_tokens=max_new_tokens,
-                        repetition_penalty=repetition_penalty,
-                        do_sample=True,
-                        temperature=temperature,
-                        top_p=top_p
-                    )
-                    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                else:
-                    response = "Model is not available. Please try again later."
-                logger.debug(f"Model response: {response}")
-                return response
+# Function for model inference
+def model_inference(
+        user_prompt,
+        chat_history,
+        web_search,
+        temperature,
+        max_new_tokens,
+        repetition_penalty,
+        top_p,
+        tokenizer  # Pass tokenizer as an argument
+):
+    logger.debug(f"Starting model inference with user prompt: {user_prompt}, chat history: {chat_history}, web_search: {web_search}")
+    if not isinstance(user_prompt, dict):
+        logger.error("Invalid input format. Expected a dictionary.")
+        return "Invalid input format. Expected a dictionary."
+
+    if "files" not in user_prompt:
+        user_prompt["files"] = []
+
+    if not user_prompt["files"]:
+        if web_search:
+            logger.debug("Performing news search")
+            news_results = fetch_news(user_prompt["text"])
+            news_text = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in news_results])
+            formatted_prompt = format_prompt(f"{user_prompt['text']} [NEWS] {news_text}", chat_history)
+        else:
+            formatted_prompt = format_prompt(user_prompt["text"], chat_history)
+
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(DEVICE)
+        if model:
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                repetition_penalty=repetition_penalty,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p
+            )
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         else:
-            return "Image input not supported in this implementation."
-    except Exception as e:
-        logger.error(f"Error during model inference: {e}")
-        return "An error occurred during model inference. Please try again."
+            response = "Model is not available. Please try again later."
+        logger.debug(f"Model response: {response}")
+        return response
+    else:
+        return "Image input not supported in this implementation."
 
 # Define Gradio interface components
 max_new_tokens = gr.Slider(
@@ -183,27 +126,71 @@ temperature = gr.Slider(
     minimum=0.0,
     maximum=2.0,
     value=0.5,
-    step=0.1,
+    step=0.05,
+    visible=True,
     interactive=True,
-    label="Temperature",
-    info="Control randomness: lower temperature produces less randomness.",
+    label="Sampling temperature",
+    info="Higher values will produce more diverse outputs.",
+)
+top_p = gr.Slider(
+    minimum=0.01,
+    maximum=0.99,
+    value=0.9,
+    step=0.01,
+    visible=True,
+    interactive=True,
+    label="Top P",
+    info="Higher values are equivalent to sampling more low-probability tokens.",
 )
-web_search = gr.Checkbox(label="Enable Web Search", default=False, description="Enable web search for better responses")
 
-# Define the Gradio interface
-gr.Interface(
-    fn=model_inference,
-    inputs=[
-        gr.Textbox(label="User Input", placeholder="Type your input here..."),
-        gr.MultiText(label="Chat History", placeholder="User: ...\nBot: ...", optional=True),
+# Create a chatbot interface
+chatbot = gr.Chatbot(
+    label="OpenGPT-4o-Chatty",
+    show_copy_button=True,
+    likeable=True,
+    layout="panel"
+)
+
+# Define Gradio interface
+def chat_interface(user_input, history, web_search, decoding_strategy, temperature, max_new_tokens, repetition_penalty, top_p):
+    logger.debug(f"Chat interface called with user_input: {user_input}")
+    if isinstance(user_input, str):
+        user_input = {"text": user_input, "files": []}
+    response = model_inference(
+        user_input,
+        history,
         web_search,
         temperature,
         max_new_tokens,
         repetition_penalty,
+        top_p,
+        tokenizer=tokenizer  # Pass tokenizer to model_inference
+    )
+    history.append((user_input["text"], response))
+    logger.debug(f"Updated chat history: {history}")
+    return history, history
+
+# Create Gradio interface
+interface = gr.Interface(
+    fn=chat_interface,
+    inputs=[
+        gr.Textbox(label="User Input"),
+        gr.State([]),
+        gr.Checkbox(label="Fetch News", value=True),
         decoding_strategy,
+        temperature,
+        max_new_tokens,
+        repetition_penalty,
+        top_p
+    ],
+    outputs=[
+        chatbot,
+        gr.State([])
     ],
-    outputs=gr.Textbox(label="AI Response"),
-    live=True,
-    title="OpenGPT 4o Demo",
-    description="An AI-powered assistant that can chat with you and provide informative responses.",
-).launch()
+    title="OpenGPT-4o-Chatty",
+    description="An AI assistant capable of insightful conversations and news fetching."
+)
+
+if __name__ == "__main__":
+    logger.debug("Launching Gradio interface")
+    interface.launch()