Update app.py

app.py

@@ -2,6 +2,10 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
 from threading import Thread
+import requests
+from bs4 import BeautifulSoup
+from urllib.parse import quote_plus
+import re
 
 MODEL_NAMES = {
     "LFM2-350M": "LiquidAI/LFM2-350M",
@@ -16,35 +20,186 @@ model_cache = {}
 def load_model(model_key):
     if model_key in model_cache:
         return model_cache[model_key]
-
     model_name = MODEL_NAMES[model_key]
     print(f"Loading {model_name}...")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         dtype=torch.float16 if device == "cuda" else torch.float32,
     ).to(device)
-
     model_cache[model_key] = (tokenizer, model)
     return tokenizer, model
 
+def search_web(query, num_results=3):
+    """
+    Search DuckDuckGo and scrape results
+    """
+    try:
+        headers = {
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+        }
+        
+        # DuckDuckGo HTML search
+        search_url = f"https://html.duckduckgo.com/html/?q={quote_plus(query)}"
+        response = requests.get(search_url, headers=headers, timeout=10)
+        soup = BeautifulSoup(response.text, 'html.parser')
+        
+        results = []
+        result_divs = soup.find_all('div', class_='result')
+        
+        for div in result_divs[:num_results]:
+            try:
+                title_tag = div.find('a', class_='result__a')
+                snippet_tag = div.find('a', class_='result__snippet')
+                
+                if title_tag and snippet_tag:
+                    title = title_tag.get_text(strip=True)
+                    url = title_tag.get('href', '')
+                    snippet = snippet_tag.get_text(strip=True)
+                    
+                    results.append({
+                        'title': title,
+                        'url': url,
+                        'snippet': snippet
+                    })
+            except Exception as e:
+                continue
+        
+        return results
+    except Exception as e:
+        return [{'title': 'Search Error', 'url': '', 'snippet': f'Could not perform search: {str(e)}'}]
 
-def chat_with_model(message, history, model_choice):
-    tokenizer, model = load_model(model_choice)
-    device = model.device
+def scrape_page_content(url, max_chars=2000):
+    """
+    Scrape main content from a webpage
+    """
+    try:
+        headers = {
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+        }
+        response = requests.get(url, headers=headers, timeout=10)
+        soup = BeautifulSoup(response.text, 'html.parser')
+        
+        # Remove script and style elements
+        for script in soup(["script", "style", "nav", "footer", "header"]):
+            script.decompose()
+        
+        # Get text
+        text = soup.get_text()
+        lines = (line.strip() for line in text.splitlines())
+        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
+        text = ' '.join(chunk for chunk in chunks if chunk)
+        
+        return text[:max_chars]
+    except Exception as e:
+        return f"Could not scrape page: {str(e)}"
+
+def detect_urls(text):
+    """
+    Detect URLs in text
+    """
+    url_pattern = r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'
+    return re.findall(url_pattern, text)
 
-    prompt = "You are LFM2, an intelligent and conversational AI assistant designed to help users with questions, problem-solving, and creative tasks. You communicate clearly, reason carefully, and explain your thoughts in an easy-to-understand way. Stay friendly, professional, and curious. If the user’s request is ambiguous, ask clarifying questions before proceeding."
+def extract_search_commands(text):
+    """
+    Extract SEARCH[query] commands from text
+    """
+    pattern = r'SEARCH\[(.*?)\]'
+    matches = re.findall(pattern, text)
+    return matches
+
+def process_with_search(message, history, tokenizer, model, enable_search):
+    """
+    Check if we need to search or scrape URLs
+    """
+    if not enable_search:
+        return message, False, None
+    
+    search_results_text = ""
+    searched = False
+    
+    # Check for URLs in the message - if found, scrape them
+    urls = detect_urls(message)
+    if urls:
+        for url in urls[:3]:  # Limit to 3 URLs
+            content = scrape_page_content(url, max_chars=3000)
+            search_results_text += f"\n\n[Content from {url}]\n{content}\n\n"
+            searched = True
+        
+        if searched:
+            enhanced_message = f"{message}\n\n{search_results_text}"
+            return enhanced_message, True, "url"
+    
+    # If no URLs, generate initial response to see if model wants to search
+    device = model.device
+    prompt = "You are LFM2, an intelligent and conversational AI assistant. You have access to web search capabilities. When you need current information or don't know something, you can use the command SEARCH[query] to search the web.\n\n"
+    
     for msg in history:
         role = msg["role"]
         content = msg["content"]
         prompt += f"{role.capitalize()}: {content}\n"
     prompt += f"User: {message}\nAssistant:"
+    
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+        )
+    
+    initial_response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    
+    # Check for search commands
+    search_queries = extract_search_commands(initial_response)
+    
+    if search_queries:
+        search_results_text = ""
+        for query in search_queries[:2]:  # Limit to 2 searches
+            results = search_web(query)
+            search_results_text += f"\n\n[Search results for: {query}]\n"
+            for i, result in enumerate(results, 1):
+                search_results_text += f"{i}. {result['title']}\n{result['snippet']}\n{result['url']}\n\n"
+        
+        # Add search results to the conversation
+        enhanced_message = f"{message}\n\n{search_results_text}"
+        return enhanced_message, True, "search"
+    
+    return message, False, None
 
+def chat_with_model(message, history, model_choice, enable_search):
+    tokenizer, model = load_model(model_choice)
+    device = model.device
+    
+    # Process search if needed
+    enhanced_message, searched, search_type = process_with_search(
+        message, history, tokenizer, model, enable_search
+    )
+    
+    # Show search indicator
+    if searched:
+        if search_type == "url":
+            search_indicator = [{"role": "assistant", "content": "🔗 Reading the provided URL..."}]
+        else:
+            search_indicator = [{"role": "assistant", "content": "🔍 Searching the web..."}]
+        yield history + [{"role": "user", "content": message}] + search_indicator
+    
+    prompt = "You are LFM2, an intelligent and conversational AI assistant designed to help users with questions, problem-solving, and creative tasks. You communicate clearly, reason carefully, and explain your thoughts in an easy-to-understand way. Stay friendly, professional, and curious. If the user's request is ambiguous, ask clarifying questions before proceeding.\n\nWEB SEARCH CAPABILITY:\n- You have access to web search. When you need current information or don't know something, use: SEARCH[your query here]\n- Only search when necessary - use your existing knowledge first\n- After searching, you'll receive results to incorporate into your response\n- Always cite sources when using search results\n\n"
+    
+    for msg in history:
+        role = msg["role"]
+        content = msg["content"]
+        prompt += f"{role.capitalize()}: {content}\n"
+    prompt += f"User: {enhanced_message}\nAssistant:"
+    
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
-
+    
     generation_kwargs = dict(
         **inputs,
         streamer=streamer,
@@ -53,52 +208,74 @@ def chat_with_model(message, history, model_choice):
         top_p=0.9,
         do_sample=True,
     )
-
+    
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
+    
     partial_text = ""
     for new_text in streamer:
         partial_text += new_text
+        # Remove SEARCH[] commands from output
+        cleaned_text = re.sub(r'SEARCH\[.*?\]', '', partial_text)
         yield history + [
             {"role": "user", "content": message},
-            {"role": "assistant", "content": partial_text},
+            {"role": "assistant", "content": cleaned_text},
         ]
 
-
 def create_demo():
-    with gr.Blocks(title="LiquidAI Chat Playground") as demo:
+    with gr.Blocks(title="LiquidAI Chat Playground", theme=gr.themes.Soft()) as demo:
         gr.Markdown("## 💧 LiquidAI Chat Playground")
-
-        model_choice = gr.Dropdown(
-            label="Select Model",
-            choices=list(MODEL_NAMES.keys()),
-            value="LFM2-1.2B"
-        )
-
+        
+        with gr.Row():
+            model_choice = gr.Dropdown(
+                label="Select Model",
+                choices=list(MODEL_NAMES.keys()),
+                value="LFM2-1.2B",
+                scale=3
+            )
+            enable_search = gr.Checkbox(
+                label="Enable Web Search",
+                value=True,
+                scale=1
+            )
+        
         chatbot = gr.Chatbot(
             label="Chat with LiquidAI",
             type="messages",
             height=450
         )
-
         msg = gr.Textbox(label="Your message", placeholder="Type something...")
         clear = gr.Button("Clear")
-
-        def add_user_message(user_message, chat_history):
-            chat_history = chat_history + [{"role": "user", "content": user_message}]
-            return "", chat_history
-
-        msg.submit(add_user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-            chat_with_model, [msg, chatbot, model_choice], chatbot
-        )
-
+        
+        def add_user_message_and_respond(user_message, chat_history, model_choice, enable_search):
+            if not user_message.strip():
+                return chat_history
+            for response in chat_with_model(user_message, chat_history, model_choice, enable_search):
+                yield response
+        
+        msg.submit(
+            add_user_message_and_respond,
+            [msg, chatbot, model_choice, enable_search],
+            chatbot
+        ).then(lambda: "", None, msg)
+        
         clear.click(lambda: [], None, chatbot, queue=False)
-
+        
+        gr.Markdown("""
+        ### Features:
+        - **Smart Search**: Model searches only when it doesn't know something
+        - **URL Access**: Paste a URL and ask the model to read it
+        - **Toggle Control**: Enable/disable web search with the checkbox
+        
+        Examples:
+        - "What's the current weather in Tokyo?" → *will search if needed*
+        - "Explain quantum physics" → *answers from knowledge*
+        - "Read this article: https://example.com" → *scrapes the URL*
+        """)
+    
     return demo
 
-
 if __name__ == "__main__":
     demo = create_demo()
     demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)