| import os |
| import time |
| import gc |
| import sys |
| import threading |
| from itertools import islice |
| from datetime import datetime |
| import re |
| from typing import List, Dict, Any, Optional, Tuple, Generator |
| from dataclasses import dataclass |
| import logging |
| import gradio as gr |
| import torch |
| from transformers import pipeline, TextIteratorStreamer |
| from transformers import AutoTokenizer |
| from bs4 import BeautifulSoup |
| import requests |
| from urllib.parse import quote_plus |
| import json |
| import urllib.parse |
| from config import MODELS |
|
|
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') |
| logger = logging.getLogger(__name__) |
|
|
| cancel_event = threading.Event() |
|
|
| ACCESS_TOKEN = os.environ.get('HF_TOKEN', '') |
| if ACCESS_TOKEN == '': |
| ACCESS_TOKEN = None |
|
|
| PIPELINES = {} |
| SEARCH_TIMEOUT_DEFAULT = 5.0 |
|
|
| @dataclass |
| class SearchResult: |
| title: str |
| snippet: str |
| url: Optional[str] = None |
| |
| def format(self, max_chars: int = 50) -> str: |
| snippet = self.snippet[:max_chars] + "..." if len(self.snippet) > max_chars else self.snippet |
| return f"{self.title} - {snippet}" |
|
|
| @dataclass |
| class GenerationConfig: |
| max_tokens: int = 1024 |
| temperature: float = 0.7 |
| top_k: int = 40 |
| top_p: float = 0.9 |
| repetition_penalty: float = 1.2 |
| |
| def to_dict(self) -> Dict[str, Any]: |
| return { |
| 'max_new_tokens': self.max_tokens, |
| 'temperature': self.temperature, |
| 'top_k': self.top_k, |
| 'top_p': self.top_p, |
| 'repetition_penalty': self.repetition_penalty, |
| } |
|
|
| class SearchEngine: |
| USER_AGENTS = [ |
| 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36', |
| 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36', |
| 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36' |
| ] |
| |
| @staticmethod |
| def _get_headers() -> Dict[str, str]: |
| return { |
| 'User-Agent': SearchEngine.USER_AGENTS[0], |
| 'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8', |
| 'Accept-Language': 'en-US,en;q=0.5', |
| 'Accept-Encoding': 'gzip, deflate', |
| 'Connection': 'keep-alive', |
| 'Upgrade-Insecure-Requests': '1', |
| 'Cache-Control': 'max-age=0' |
| } |
|
|
| class GoogleSearch(SearchEngine): |
| @staticmethod |
| def search(query: str, max_results: int = 6, max_chars: int = 50) -> List[SearchResult]: |
| encoded_query = quote_plus(query) |
| search_urls = [ |
| f"https://www.google.com/search?q={encoded_query}&safe=off&num={max_results}", |
| f"https://www.google.com/search?q={encoded_query}&safe=off&num={max_results}&hl=en", |
| f"https://www.google.com/webhp?safe=off&q={encoded_query}&num={max_results}" |
| ] |
| |
| for user_agent in SearchEngine.USER_AGENTS: |
| headers = SearchEngine._get_headers() |
| headers['User-Agent'] = user_agent |
| |
| for search_url in search_urls: |
| try: |
| response = requests.get(search_url, headers=headers, timeout=15, verify=True) |
| response.raise_for_status() |
| |
| soup = BeautifulSoup(response.text, 'html.parser') |
| |
| selectors = [ |
| ('div', 'g'), |
| ('div', 'tF2Cxc'), |
| ('div', 'MjjYud'), |
| ('div', 'yuRUbf') |
| ] |
| |
| search_results = [] |
| for tag, class_name in selectors: |
| search_results = soup.find_all(tag, class_=class_name) |
| if search_results: |
| break |
| |
| if not search_results: |
| search_results = soup.find_all('div', class_=re.compile(r'^(g|tF2Cxc|MjjYud|yuRUbf)')) |
| |
| results = [] |
| for result in search_results[:max_results]: |
| try: |
| title_elem = result.find('h3') or result.find('h2') |
| if not title_elem: |
| continue |
| |
| snippet_elem = result.find('div', class_='VwiC3b') or \ |
| result.find('div', class_='IsZvec') or \ |
| result.find('div', class_='lEBKkf') |
| |
| link_elem = result.find('a') |
| if not link_elem: |
| continue |
| |
| link = link_elem.get('href', '') |
| if link.startswith('/url?q='): |
| link = urllib.parse.unquote(link.split('/url?q=')[1].split('&')[0]) |
| |
| if not link.startswith('http'): |
| continue |
| |
| title = title_elem.text.strip() |
| snippet = snippet_elem.text.strip() if snippet_elem else "" |
| snippet = ' '.join(snippet.split()) |
| |
| if title and snippet: |
| results.append(SearchResult(title=title, snippet=snippet, url=link)) |
| |
| except Exception as e: |
| logger.debug(f"Error parsing Google result: {e}") |
| continue |
| |
| if results: |
| return results |
| |
| except Exception as e: |
| logger.debug(f"Google search attempt failed: {e}") |
| continue |
| |
| return [] |
|
|
| class DuckDuckGoSearch(SearchEngine): |
| @staticmethod |
| def search(query: str, max_results: int = 6, max_chars: int = 50) -> List[SearchResult]: |
| try: |
| from ddgs import DDGS |
| with DDGS() as ddgs: |
| results = [] |
| for r in islice(ddgs.text(query, region="wt-wt", safesearch="off", timelimit="y"), max_results): |
| title = r.get('title', 'No Title') |
| body = r.get('body', '') |
| results.append(SearchResult(title=title, snippet=body)) |
| return results |
| except Exception as e: |
| logger.debug(f"DuckDuckGo search failed: {e}") |
| return [] |
|
|
| class BingSearch(SearchEngine): |
| @staticmethod |
| def search(query: str, max_results: int = 6, max_chars: int = 50) -> List[SearchResult]: |
| try: |
| headers = SearchEngine._get_headers() |
| search_url = f"https://www.bing.com/search?q={quote_plus(query)}&safeSearch=off&count={max_results}" |
| |
| response = requests.get(search_url, headers=headers, timeout=10) |
| response.raise_for_status() |
| |
| soup = BeautifulSoup(response.text, 'html.parser') |
| results = [] |
| |
| for result in soup.find_all('li', class_='b_algo')[:max_results]: |
| try: |
| title_elem = result.find('h2') |
| snippet_elem = result.find('p') |
| |
| if title_elem and snippet_elem: |
| title = title_elem.text.strip() |
| snippet = snippet_elem.text.strip() |
| results.append(SearchResult(title=title, snippet=snippet)) |
| |
| except Exception as e: |
| logger.debug(f"Error parsing Bing result: {e}") |
| continue |
| |
| return results |
| except Exception as e: |
| logger.debug(f"Bing search failed: {e}") |
| return [] |
|
|
| class SearchManager: |
| _engines = [GoogleSearch, DuckDuckGoSearch, BingSearch] |
| |
| @classmethod |
| def search(cls, query: str, max_results: int = 6, max_chars: int = 50, timeout: float = 5.0) -> List[SearchResult]: |
| for engine_cls in cls._engines: |
| try: |
| result_container = [] |
| search_thread = threading.Thread( |
| target=lambda: result_container.extend(engine_cls.search(query, max_results, max_chars)) |
| ) |
| search_thread.daemon = True |
| search_thread.start() |
| search_thread.join(timeout=timeout) |
| |
| if result_container: |
| logger.info(f"Search successful with {engine_cls.__name__}: {len(result_container)} results") |
| return result_container |
| |
| except Exception as e: |
| logger.warning(f"Search engine {engine_cls.__name__} failed: {e}") |
| continue |
| |
| return [] |
|
|
| class ModelManager: |
| _pipelines = {} |
| _lock = threading.Lock() |
| |
| @classmethod |
| def load_pipeline(cls, model_name: str) -> pipeline: |
| with cls._lock: |
| if model_name in cls._pipelines: |
| return cls._pipelines[model_name] |
| |
| repo = MODELS[model_name]["repo_id"] |
| |
| try: |
| tokenizer = AutoTokenizer.from_pretrained( |
| repo, |
| token=ACCESS_TOKEN if ACCESS_TOKEN else None |
| ) |
| except Exception as e: |
| logger.warning(f"Failed to load tokenizer with token, trying without: {e}") |
| tokenizer = AutoTokenizer.from_pretrained(repo) |
| |
| for dtype in (torch.bfloat16, torch.float16, torch.float32): |
| try: |
| pipe_kwargs = { |
| 'task': "text-generation", |
| 'model': repo, |
| 'tokenizer': tokenizer, |
| 'trust_remote_code': True, |
| 'dtype': dtype, |
| 'device_map': "auto", |
| 'use_cache': True, |
| } |
| if ACCESS_TOKEN: |
| pipe_kwargs['token'] = ACCESS_TOKEN |
| |
| pipe = pipeline(**pipe_kwargs) |
| cls._pipelines[model_name] = pipe |
| return pipe |
| except Exception as e: |
| logger.warning(f"Failed to load with {dtype}: {e}") |
| continue |
| |
| pipe_kwargs = { |
| 'task': "text-generation", |
| 'model': repo, |
| 'tokenizer': tokenizer, |
| 'trust_remote_code': True, |
| 'device_map': "auto", |
| 'use_cache': True, |
| } |
| if ACCESS_TOKEN: |
| pipe_kwargs['token'] = ACCESS_TOKEN |
| |
| pipe = pipeline(**pipe_kwargs) |
| cls._pipelines[model_name] = pipe |
| return pipe |
|
|
| class PromptBuilder: |
| @staticmethod |
| def format_conversation(history: List[Dict], system_prompt: str, tokenizer) -> str: |
| if hasattr(tokenizer, "chat_template") and tokenizer.chat_template: |
| messages = [{"role": "system", "content": system_prompt.strip()}] + history |
| return tokenizer.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=True |
| ) |
| else: |
| prompt = f"{system_prompt.strip()}\n" |
| for msg in history: |
| if msg['role'] == 'user': |
| prompt += f"User: {msg['content'].strip()}\n" |
| elif msg['role'] == 'assistant': |
| prompt += f"Assistant: {msg['content'].strip()}\n" |
| |
| if not prompt.strip().endswith("Assistant:"): |
| prompt += "Assistant: " |
| return prompt |
| |
| @staticmethod |
| def build_search_context(search_results: List[SearchResult], system_prompt: str, user_query: str) -> str: |
| if not search_results: |
| return system_prompt.strip() |
| |
| formatted_results = "\n".join(f"[{i+1}] {r.format()}" for i, r in enumerate(search_results)) |
| |
| return f"""{system_prompt.strip()} |
| |
| # SEARCH CONTEXT (TRUSTED SOURCES ONLY) |
| Below are search results. Treat them as the ONLY source of truth for answering. |
| {formatted_results} |
| |
| RULES (VERY IMPORTANT): |
| - Do NOT use outside knowledge. Do NOT guess or fill missing information. |
| - If the answer is not clearly supported by the search results, say: "Not enough information in the provided sources." |
| - Every factual statement must be directly supported by at least one citation [citation:X]. |
| - Do NOT add explanations, examples, or background that are not explicitly present in the sources. |
| - Do NOT paraphrase beyond what is necessary for clarity. |
| - If sources conflict, mention the conflict and cite both. |
| - If multiple sources are used, distribute citations per sentence, not only at the end. |
| |
| CITATION RULES: |
| - Use inline citations like this: [citation:1] |
| - If multiple sources support a sentence: [citation:1][citation:3] |
| - Never place all citations only at the end. |
| |
| ANSWER POLICY: |
| - Be concise and strictly grounded. |
| - No speculation, no assumptions, no "likely", no "probably". |
| - If the user requests a list, only include items explicitly found in sources. |
| - If sources are insufficient, stop and ask for more data instead of guessing. |
| |
| DATE CONTEXT: |
| - Today is {datetime.now().strftime('%Y-%m-%d')} (use only for time reference, not for assumptions). |
| |
| USER QUESTION: |
| {user_query}""" |
|
|
| class StreamProcessor: |
| @staticmethod |
| def process_stream(streamer: TextIteratorStreamer, history: List[Dict]) -> Generator[Tuple[List[Dict], str], None, None]: |
| thought_buf = '' |
| answer_buf = '' |
| in_thought = False |
| assistant_message_started = False |
| |
| for chunk in streamer: |
| if cancel_event.is_set(): |
| if assistant_message_started and history and history[-1]['role'] == 'assistant': |
| history[-1]['content'] += " [Generation Canceled]" |
| yield history, "Generation canceled by user." |
| break |
| |
| text = chunk |
| |
| if not in_thought and '<think>' in text: |
| in_thought = True |
| history.append({'role': 'assistant', 'content': '', 'metadata': {'title': '💭 Thought'}}) |
| assistant_message_started = True |
| after = text.split('<think>', 1)[1] |
| thought_buf += after |
| |
| if '</think>' in thought_buf: |
| before, after2 = thought_buf.split('</think>', 1) |
| history[-1]['content'] = before.strip() |
| in_thought = False |
| answer_buf = after2 |
| history.append({'role': 'assistant', 'content': answer_buf}) |
| else: |
| history[-1]['content'] = thought_buf |
| yield history, "" |
| continue |
| |
| if in_thought: |
| thought_buf += text |
| if '</think>' in thought_buf: |
| before, after2 = thought_buf.split('</think>', 1) |
| history[-1]['content'] = before.strip() |
| in_thought = False |
| answer_buf = after2 |
| history.append({'role': 'assistant', 'content': answer_buf}) |
| else: |
| history[-1]['content'] = thought_buf |
| yield history, "" |
| continue |
| |
| if not assistant_message_started: |
| history.append({'role': 'assistant', 'content': ''}) |
| assistant_message_started = True |
| |
| answer_buf += text |
| history[-1]['content'] = answer_buf.strip() |
| yield history, "" |
|
|
| def chat_response( |
| user_msg: str, |
| chat_history: List[Dict], |
| system_prompt: str, |
| enable_search: bool, |
| max_results: int, |
| max_chars: int, |
| model_name: str, |
| max_tokens: int, |
| temperature: float, |
| top_k: int, |
| top_p: float, |
| repeat_penalty: float, |
| search_timeout: float |
| ) -> Generator[Tuple[List[Dict], str], None, None]: |
| cancel_event.clear() |
| history = list(chat_history or []) |
| history.append({'role': 'user', 'content': user_msg}) |
| |
| search_results: List[SearchResult] = [] |
| search_debug = "Web search disabled." |
| |
| if enable_search: |
| search_debug = "🔍 Searching across multiple engines..." |
| try: |
| search_results = SearchManager.search( |
| user_msg, |
| int(max_results), |
| int(max_chars), |
| float(search_timeout) |
| ) |
| |
| if search_results: |
| search_debug = f"✅ Search completed - Found {len(search_results)} results\n\n" + "\n".join( |
| f"- {r.format(int(max_chars))}" for r in search_results |
| ) |
| else: |
| search_debug = "❌ No search results found. Check internet connection or try again." |
| except Exception as e: |
| search_debug = f"❌ Search failed: {str(e)}" |
| logger.error(f"Search error: {e}") |
| |
| try: |
| if enable_search and search_results: |
| enriched_prompt = PromptBuilder.build_search_context( |
| search_results, |
| system_prompt, |
| user_msg |
| ) |
| else: |
| enriched_prompt = system_prompt.strip() |
| |
| pipe = ModelManager.load_pipeline(model_name) |
| |
| prompt = PromptBuilder.format_conversation(history, enriched_prompt, pipe.tokenizer) |
| prompt_debug = f"\n\n--- Prompt Preview ---\n```\n{prompt[:500]}...\n```" if len(prompt) > 500 else f"\n\n--- Prompt Preview ---\n```\n{prompt}\n```" |
| |
| config = GenerationConfig( |
| max_tokens=max_tokens, |
| temperature=temperature, |
| top_k=top_k, |
| top_p=top_p, |
| repetition_penalty=repeat_penalty |
| ) |
| |
| streamer = TextIteratorStreamer( |
| pipe.tokenizer, |
| skip_prompt=True, |
| skip_special_tokens=True |
| ) |
| |
| gen_kwargs = config.to_dict() |
| gen_kwargs['streamer'] = streamer |
| gen_kwargs['return_full_text'] = False |
| |
| gen_thread = threading.Thread( |
| target=pipe, |
| args=(prompt,), |
| kwargs=gen_kwargs |
| ) |
| gen_thread.start() |
| |
| yield history, search_debug |
| |
| for history_update, debug_update in StreamProcessor.process_stream(streamer, history): |
| yield history_update, debug_update |
| |
| gen_thread.join(timeout=5.0) |
| yield history, search_debug + prompt_debug |
| |
| except GeneratorExit: |
| logger.info("Generation cancelled by user") |
| return |
| except Exception as e: |
| logger.error(f"Generation error: {e}") |
| history.append({'role': 'assistant', 'content': f"Error: {str(e)}"}) |
| yield history, search_debug |
| finally: |
| gc.collect() |
|
|
| def get_model_size(model_name: str) -> float: |
| return MODELS.get(model_name, {}).get("params_b", 4.0) |
|
|
| def get_duration_estimate( |
| model_name: str, |
| enable_search: bool, |
| max_tokens: int, |
| search_timeout: float |
| ) -> float: |
| model_size = get_model_size(model_name) |
| use_aot = model_size >= 2 |
| |
| base_duration = 20 if not use_aot else 40 |
| token_duration = max_tokens * 0.005 |
| search_duration = 10 if enable_search else 0 |
| aot_compilation = 20 if use_aot else 0 |
| |
| return base_duration + token_duration + search_duration + aot_compilation |
|
|
| def update_duration_estimate( |
| model_name: str, |
| enable_search: bool, |
| max_results: int, |
| max_chars: int, |
| max_tokens: int, |
| search_timeout: float |
| ) -> str: |
| try: |
| duration = get_duration_estimate(model_name, enable_search, max_tokens, search_timeout) |
| model_size = get_model_size(model_name) |
| |
| return f"""⏱️ **Estimated GPU Time: {duration:.1f} seconds** |
| |
| 📊 **Model Size:** {model_size:.1f}B parameters |
| 🔍 **Web Search:** {'Enabled (Multi-Engine)' if enable_search else 'Disabled'}""" |
| except Exception as e: |
| logger.error(f"Error calculating estimate: {e}") |
| return f"⚠️ Error calculating estimate: {e}" |
|
|
| def update_default_prompt(enable_search: bool) -> str: |
| return "You are a helpful assistant." |
|
|
| with gr.Blocks( |
| title="LLM Inference", |
| theme=gr.themes.Soft( |
| primary_hue="blue", |
| secondary_hue="blue", |
| neutral_hue="slate", |
| radius_size="lg", |
| font=[gr.themes.GoogleFont("Syne"), "Arial", "sans-serif"] |
| ), |
| css=""" |
| .duration-estimate { background: linear-gradient(135deg, #667eea15 0%, #764ba215 100%); border-left: 4px solid #667eea; padding: 12px; border-radius: 8px; margin: 16px 0; } |
| .chatbot { border-radius: 12px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); } |
| button.primary { font-weight: 600; } |
| .gradio-accordion { margin-bottom: 12px; } |
| """ |
| ) as demo: |
| gr.Markdown(""" |
| # 🧠 LLM Inference with Multi-Engine Search |
| """) |
| |
| with gr.Row(): |
| with gr.Column(scale=3): |
| with gr.Group(): |
| gr.Markdown("### ⚙️ Core Settings") |
| model_dd = gr.Dropdown( |
| label="🤖 Model", |
| choices=list(MODELS.keys()), |
| value="Qwen3-1.7B", |
| info="Select the language model to use" |
| ) |
| search_chk = gr.Checkbox( |
| label="🔍 Enable Web Search", |
| value=False, |
| info="Search across Google, DuckDuckGo, and Bing (no API required)" |
| ) |
| sys_prompt = gr.Textbox(label="📝 System Prompt", lines=3, value=update_default_prompt(False), placeholder="Define the assistant's behavior and personality...") |
| |
| duration_display = gr.Markdown( |
| value=update_duration_estimate("Qwen3-1.7B", False, 4, 50, 1024, 5.0), |
| elem_classes="duration-estimate" |
| ) |
| |
| with gr.Accordion("🎛️ Advanced Generation Parameters", open=False): |
| max_tok = gr.Slider( |
| 64, 16384, value=1024, step=32, |
| label="Max Tokens", |
| info="Maximum length of generated response" |
| ) |
| temp = gr.Slider( |
| 0.1, 2.0, value=0.7, step=0.1, |
| label="Temperature", |
| info="Higher = more creative, Lower = more focused" |
| ) |
| with gr.Row(): |
| k = gr.Slider( |
| 1, 100, value=40, step=1, |
| label="Top-K", |
| info="Number of top tokens to consider" |
| ) |
| p = gr.Slider( |
| 0.1, 1.0, value=0.9, step=0.05, |
| label="Top-P", |
| info="Nucleus sampling threshold" |
| ) |
| rp = gr.Slider( |
| 1.0, 2.0, value=1.2, step=0.1, |
| label="Repetition Penalty", |
| info="Penalize repeated tokens" |
| ) |
| |
| with gr.Accordion("🌐 Web Search Settings", open=False, visible=False) as search_settings: |
| mr = gr.Number( |
| value=4, precision=0, |
| label="Max Results", |
| info="Number of search results to retrieve" |
| ) |
| mc = gr.Number( |
| value=50, precision=0, |
| label="Max Chars/Result", |
| info="Character limit per search result" |
| ) |
| st = gr.Slider( |
| minimum=0.0, maximum=30.0, step=0.5, value=5.0, |
| label="Search Timeout (s)", |
| info="Maximum time to wait for search results" |
| ) |
| gr.Markdown(""" |
| ⚠️ **Search Engines:** |
| - Google (primary) |
| - DuckDuckGo (fallback) |
| - Bing (fallback) |
| |
| SafeSearch is **OFF** for comprehensive results. |
| """) |
| |
| with gr.Row(): |
| clr = gr.Button("🗑️ Clear Chat", variant="secondary", scale=1) |
| |
| with gr.Column(scale=7): |
| chat = gr.Chatbot( |
| type="messages", |
| height=600, |
| label="💬 Conversation", |
| show_copy_button=True, |
| avatar_images=( |
| "data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='40' height='40'%3E%3Crect width='40' height='40' rx='20' fill='%23f093fb'/%3E%3Ctext x='20' y='28' text-anchor='middle' font-size='20' fill='white' font-family='Arial'%3E👤%3C/text%3E%3C/svg%3E", |
| "data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='40' height='40'%3E%3Crect width='40' height='40' rx='20' fill='%23667eea'/%3E%3Ctext x='20' y='28' text-anchor='middle' font-size='20' fill='white' font-family='Arial'%3E🤖%3C/text%3E%3C/svg%3E" |
| ), |
| bubble_full_width=False, |
| render_markdown=True, |
| sanitize_html=False |
| ) |
| |
| with gr.Row(): |
| txt = gr.Textbox( |
| placeholder="💭 Type your message here... (Press Enter to send)", |
| scale=9, |
| container=False, |
| show_label=False, |
| lines=1, |
| max_lines=5 |
| ) |
| with gr.Column(scale=1, min_width=120): |
| submit_btn = gr.Button("📤 Send", variant="primary", size="lg") |
| cancel_btn = gr.Button("⏹️ Stop", variant="stop", visible=False, size="lg") |
| |
| gr.Examples( |
| examples=[ |
| ["Explain quantum computing in simple terms"], |
| ["Write a Python function to calculate fibonacci numbers"], |
| ["What are the latest developments in AI? (Enable web search)"], |
| ["Tell me a creative story about a time traveler"], |
| ["Help me debug this code: def add(a,b): return a+b+1"] |
| ], |
| inputs=txt, |
| label="💡 Example Prompts" |
| ) |
| |
| with gr.Accordion("🔍 Debug Info", open=False): |
| dbg = gr.Markdown() |
| |
| gr.Markdown(""" |
| --- |
| 💡 **Tips:** |
| - Use **Advanced Parameters** to fine-tune creativity and response length |
| - Enable **Web Search** for real-time information (uses multiple search engines) |
| - SafeSearch is **OFF** for comprehensive results |
| - Try different **models** for various tasks (reasoning, coding, general chat) |
| - Click the **Copy** button on responses to save them to your clipboard |
| """, elem_classes="footer") |
|
|
| chat_inputs = [txt, chat, sys_prompt, search_chk, mr, mc, model_dd, max_tok, temp, k, p, rp, st] |
| ui_components = [chat, dbg, txt, submit_btn, cancel_btn] |
|
|
| def submit_and_manage_ui(user_msg, chat_history, *args): |
| if not user_msg.strip(): |
| yield {} |
| return |
|
|
| yield { |
| txt: gr.update(value="", interactive=False), |
| submit_btn: gr.update(interactive=False), |
| cancel_btn: gr.update(visible=True), |
| } |
|
|
| cancelled = False |
| try: |
| backend_args = [user_msg, chat_history] + list(args) |
| for response_chunk in chat_response(*backend_args): |
| yield { |
| chat: response_chunk[0], |
| dbg: response_chunk[1], |
| } |
| except GeneratorExit: |
| cancelled = True |
| print("Generation cancelled by user.") |
| raise |
| except Exception as e: |
| print(f"An error occurred during generation: {e}") |
| error_history = (chat_history or []) + [ |
| {'role': 'user', 'content': user_msg}, |
| {'role': 'assistant', 'content': f"**An error occurred:** {str(e)}"} |
| ] |
| yield {chat: error_history} |
| finally: |
| if not cancelled: |
| print("Resetting UI state.") |
| yield { |
| txt: gr.update(interactive=True), |
| submit_btn: gr.update(interactive=True), |
| cancel_btn: gr.update(visible=False), |
| } |
|
|
| def set_cancel_flag(): |
| cancel_event.set() |
| print("Cancellation signal sent.") |
| |
| def reset_ui_after_cancel(): |
| cancel_event.clear() |
| print("UI reset after cancellation.") |
| return { |
| txt: gr.update(interactive=True), |
| submit_btn: gr.update(interactive=True), |
| cancel_btn: gr.update(visible=False), |
| } |
|
|
| submit_event = txt.submit( |
| fn=submit_and_manage_ui, |
| inputs=chat_inputs, |
| outputs=ui_components, |
| ) |
| submit_btn.click( |
| fn=submit_and_manage_ui, |
| inputs=chat_inputs, |
| outputs=ui_components, |
| ) |
|
|
| cancel_btn.click( |
| fn=set_cancel_flag, |
| cancels=[submit_event] |
| ).then( |
| fn=reset_ui_after_cancel, |
| outputs=ui_components |
| ) |
|
|
| duration_inputs = [model_dd, search_chk, mr, mc, max_tok, st] |
| for component in duration_inputs: |
| component.change(fn=update_duration_estimate, inputs=duration_inputs, outputs=duration_display) |
|
|
| def toggle_search_settings(enabled): |
| return gr.update(visible=enabled) |
| |
| search_chk.change( |
| fn=lambda enabled: (update_default_prompt(enabled), gr.update(visible=enabled)), |
| inputs=search_chk, |
| outputs=[sys_prompt, search_settings] |
| ) |
| |
| clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg]) |
| |
| demo.launch(share=True) |
