Update app.py

app.py

@@ -1,161 +1,155 @@
-import os
-import urllib
-import requests
-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
 
-# Set up logging
+# Initialize logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# 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"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 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 perform a Google search and return the results
-def search(term, num_results=2, lang="en", timeout=5, safe="active", ssl_verify=None):
-    logger.debug(f"Starting search for term: {term}")
-    escaped_term = urllib.parse.quote_plus(term)
-    start = 0
-    all_results = []
-    max_chars_per_page = 8000
+# 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'
+]
 
-    with requests.Session() as session:
-        while start < num_results:
-            try:
-                resp = session.get(
-                    url="https://www.google.com/search",
-                    headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
-                    params={
-                        "q": term,
-                        "num": num_results - start,
-                        "hl": lang,
-                        "start": start,
-                        "safe": safe,
-                    },
-                    timeout=timeout,
-                    verify=ssl_verify,
-                )
-                resp.raise_for_status()
-                soup = BeautifulSoup(resp.text, "html.parser")
-                result_block = soup.find_all("div", attrs={"class": "g"})
-                if not result_block:
-                    start += 1
-                    continue
-                for result in result_block:
-                    link = result.find("a", href=True)
-                    if link:
-                        link = link["href"]
-                        try:
-                            webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"})
-                            webpage.raise_for_status()
-                            visible_text = extract_text_from_webpage(webpage.text)
-                            if len(visible_text) > max_chars_per_page:
-                                visible_text = visible_text[:max_chars_per_page] + "..."
-                            all_results.append({"link": link, "text": visible_text})
-                        except requests.exceptions.RequestException as e:
-                            logger.error(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)
-            except Exception as e:
-                logger.error(f"Error during search: {e}")
-                break
-    logger.debug(f"Search results: {all_results}")
-    return all_results
+def get_useragent():
+    """Returns a random user agent from the list."""
+    return random.choice(_useragent_list)
 
-# Function to extract visible text from HTML content
 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
 
-# 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 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 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 = []
 
-# 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."
+    # 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 "files" not in user_prompt:
-        user_prompt["files"] = []
+        # If no results, continue to the next batch
+        if not result_block:
+            start += 1
+            continue
 
-    if not user_prompt["files"]:
-        if web_search:
-            logger.debug("Performing web search")
-            web_results = search(user_prompt["text"])
-            web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
-            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)
+        # 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:
-                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)
+                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/Mixtral-8x7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to(DEVICE)
+
+def format_prompt(user_input, chat_history):
+    prompt = ""
+    for user, bot in chat_history:
+        prompt += f"User: {user}\nBot: {bot}\n"
+    prompt += f"User: {user_input}\nBot: "
+    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:
-                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."
+                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
+        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."
 
 # Define Gradio interface components
 max_new_tokens = gr.Slider(
@@ -189,71 +183,27 @@ temperature = gr.Slider(
     minimum=0.0,
     maximum=2.0,
     value=0.5,
-    step=0.05,
-    visible=True,
+    step=0.1,
     interactive=True,
-    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.",
-)
-
-# Create a chatbot interface
-chatbot = gr.Chatbot(
-    label="OpenGPT-4o-Chatty",
-    show_copy_button=True,
-    likeable=True,
-    layout="panel"
+    label="Temperature",
+    info="Control randomness: lower temperature produces less randomness.",
 )
+web_search = gr.Checkbox(label="Enable Web Search", default=False, description="Enable web search for better responses")
 
-# 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,
+# 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),
         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="Web Search", value=True),
         decoding_strategy,
-        temperature,
-        max_new_tokens,
-        repetition_penalty,
-        top_p
     ],
-    outputs=[
-        chatbot,
-        gr.State([])
-    ],
-    title="OpenGPT-4o-Chatty",
-    description="An AI assistant capable of insightful conversations and web search."
-)
-
-if __name__ == "__main__":
-    logger.debug("Launching Gradio interface")
-    interface.launch()
+    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()