import os import time import gc import threading from itertools import islice from datetime import datetime import gradio as gr import torch from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer from duckduckgo_search import DDGS import spaces # Import spaces early to enable ZeroGPU support # Optional: Disable GPU visibility if you wish to force CPU usage # os.environ["CUDA_VISIBLE_DEVICES"] = "" # ------------------------------ # Global Cancellation Event # ------------------------------ cancel_event = threading.Event() # ------------------------------ # Torch-Compatible Model Definitions with Adjusted Descriptions # ------------------------------ MODELS = { "Gemma-3-4B-IT": { "repo_id": "unsloth/gemma-3-4b-it", "description": "Gemma-3-4B-IT" }, "Llama-3.2-Taiwan-3B-Instruct": { "repo_id": "lianghsun/Llama-3.2-Taiwan-3B-Instruct", "description": "Llama-3.2-Taiwan-3B-Instruct" }, "MiniCPM3-4B": { "repo_id": "openbmb/MiniCPM3-4B", "description": "MiniCPM3-4B" }, "Qwen2.5-3B-Instruct": { "repo_id": "Qwen/Qwen2.5-3B-Instruct", "description": "Qwen2.5-3B-Instruct" }, "Qwen2.5-7B-Instruct": { "repo_id": "Qwen/Qwen2.5-7B-Instruct", "description": "Qwen2.5-7B-Instruct" }, "Phi-4-mini-Instruct": { "repo_id": "unsloth/Phi-4-mini-instruct", "description": "Phi-4-mini-Instruct" }, "Meta-Llama-3.1-8B-Instruct": { "repo_id": "MaziyarPanahi/Meta-Llama-3.1-8B-Instruct", "description": "Meta-Llama-3.1-8B-Instruct" }, "DeepSeek-R1-Distill-Llama-8B": { "repo_id": "unsloth/DeepSeek-R1-Distill-Llama-8B", "description": "DeepSeek-R1-Distill-Llama-8B" }, "Mistral-7B-Instruct-v0.3": { "repo_id": "MaziyarPanahi/Mistral-7B-Instruct-v0.3", "description": "Mistral-7B-Instruct-v0.3" }, "Qwen2.5-Coder-7B-Instruct": { "repo_id": "Qwen/Qwen2.5-Coder-7B-Instruct", "description": "Qwen2.5-Coder-7B-Instruct" }, } # Global cache for pipelines to avoid re-loading. PIPELINES = {} def load_pipeline(model_name): """ Load and cache a transformers pipeline for chat/text-generation. Uses the model's repo_id from MODELS and caches the pipeline for future use. """ global PIPELINES if model_name in PIPELINES: return PIPELINES[model_name] selected_model = MODELS[model_name] # Create a chat-style text-generation pipeline. pipe = pipeline( task="text-generation", model=selected_model["repo_id"], tokenizer=selected_model["repo_id"], trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto" ) PIPELINES[model_name] = pipe return pipe def retrieve_context(query, max_results=6, max_chars_per_result=600): """ Retrieve recent web search context for the given query using DuckDuckGo. Returns a formatted string with search results. """ try: with DDGS() as ddgs: results = list(islice(ddgs.text(query, region="wt-wt", safesearch="off", timelimit="y"), max_results)) context = "" for i, result in enumerate(results, start=1): title = result.get("title", "No Title") snippet = result.get("body", "")[:max_chars_per_result] context += f"Result {i}:\nTitle: {title}\nSnippet: {snippet}\n\n" return context.strip() except Exception: return "" # ---------------------------------------------------------------------------- # NEW HELPER FUNCTION: Format Conversation History into a Clean Prompt # ---------------------------------------------------------------------------- def format_conversation(conversation, system_prompt): """ Converts a list of conversation messages (each a dict with 'role' and 'content') and a system prompt into a single plain text string. This prevents raw role labels from being passed to the model. """ # Start with the system prompt. prompt = system_prompt.strip() + "\n" # Loop through conversation and format user and assistant messages. for msg in conversation: if msg["role"] == "user": prompt += "User: " + msg["content"].strip() + "\n" elif msg["role"] == "assistant": prompt += "Assistant: " + msg["content"].strip() + "\n" elif msg["role"] == "system": prompt += msg["content"].strip() + "\n" # Append the assistant cue to indicate the start of the reply. if not prompt.strip().endswith("Assistant:"): prompt += "Assistant: " return prompt # ------------------------------ # Chat Response Generation with ZeroGPU using Pipeline (Streaming Token-by-Token) # ------------------------------ @spaces.GPU(duration=60) def chat_response(user_message, chat_history, system_prompt, enable_search, max_results, max_chars, model_name, max_tokens, temperature, top_k, top_p, repeat_penalty): """ Generate a chat response by utilizing a transformers pipeline with streaming. - Appends the user's message to the conversation history. - Optionally retrieves web search context and inserts it as an additional system message. - Converts the conversation into a formatted prompt to avoid leaking role labels. - Uses the cached pipelineā€™s underlying model and tokenizer with a streamer to yield tokens as they are generated. - Yields updated conversation history token by token. """ cancel_event.clear() # Build conversation list from chat history. conversation = list(chat_history) if chat_history else [] conversation.append({"role": "user", "content": user_message}) # Retrieve web search context if enabled. debug_message = "" if enable_search: debug_message = "Initiating web search..." yield conversation, debug_message search_result = [""] def do_search(): search_result[0] = retrieve_context(user_message, max_results, max_chars) search_thread = threading.Thread(target=do_search) search_thread.start() search_thread.join(timeout=2) retrieved_context = search_result[0] if retrieved_context: debug_message = f"Web search results:\n\n{retrieved_context}" # Insert the search context as a system-level message immediately after the original system prompt. conversation.insert(1, {"role": "system", "content": f"Web search context:\n{retrieved_context}"}) else: debug_message = "Web search returned no results or timed out." else: debug_message = "Web search disabled." # Append a placeholder for the assistant's response. conversation.append({"role": "assistant", "content": ""}) try: # Format the entire conversation into a single prompt. prompt_text = format_conversation(conversation, system_prompt) # Load the pipeline. pipe = load_pipeline(model_name) # Obtain the underlying tokenizer and model. tokenizer = pipe.tokenizer model = pipe.model # Tokenize the formatted prompt. model_inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device) # Set up a streamer for token-by-token generation. streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) # Run generate in a background thread with the streamer. gen_kwargs = { "input_ids": model_inputs.input_ids, "attention_mask": model_inputs.attention_mask, "max_new_tokens": max_tokens, "temperature": temperature, "top_k": top_k, "top_p": top_p, "repetition_penalty": repeat_penalty, "streamer": streamer } thread = threading.Thread(target=model.generate, kwargs=gen_kwargs) thread.start() # Collect tokens from the streamer as they are generated. assistant_text = "" for new_text in streamer: assistant_text += new_text conversation[-1]["content"] = assistant_text yield conversation, debug_message # Update UI token by token thread.join() except Exception as e: conversation[-1]["content"] = f"Error: {e}" yield conversation, debug_message finally: gc.collect() # ------------------------------ # Cancel Function # ------------------------------ def cancel_generation(): cancel_event.set() return "Cancellation requested." # ------------------------------ # Helper Function for Default Prompt Update # ------------------------------ def update_default_prompt(enable_search): today = datetime.now().strftime('%Y-%m-%d') if enable_search: return f"You are a helpful assistant. Today is {today}. Please leverage the latest web data when responding to queries." else: return f"You are a helpful assistant. Today is {today}." # ------------------------------ # Gradio UI Definition # ------------------------------ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo: gr.Markdown("## šŸ§  ZeroGPU LLM Inference with Web Search") gr.Markdown("Interact with the model. Select your model, set your system prompt, and adjust parameters on the left.") with gr.Row(): with gr.Column(scale=3): default_model = list(MODELS.keys())[0] if MODELS else "No models available" model_dropdown = gr.Dropdown( label="Select Model", choices=list(MODELS.keys()) if MODELS else [], value=default_model, info="Choose from available models." ) # Create the Enable Web Search checkbox. enable_search_checkbox = gr.Checkbox(label="Enable Web Search", value=True, info="Include recent search context to improve answers.") # Create the System Prompt textbox with an initial value. system_prompt_text = gr.Textbox(label="System Prompt", value=update_default_prompt(enable_search_checkbox.value), lines=3, info="Define the base context for the AI's responses.") gr.Markdown("### Generation Parameters") max_tokens_slider = gr.Slider(label="Max Tokens", minimum=64, maximum=1024, value=1024, step=32, info="Maximum tokens for the response.") temperature_slider = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7, step=0.1, info="Controls the randomness of the output.") top_k_slider = gr.Slider(label="Top-K", minimum=1, maximum=100, value=40, step=1, info="Limits token candidates to the top-k tokens.") top_p_slider = gr.Slider(label="Top-P (Nucleus Sampling)", minimum=0.1, maximum=1.0, value=0.95, step=0.05, info="Limits token candidates to a cumulative probability threshold.") repeat_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.1, step=0.1, info="Penalizes token repetition to improve diversity.") gr.Markdown("### Web Search Settings") max_results_number = gr.Number(label="Max Search Results", value=10, precision=0, info="Maximum number of search results to retrieve.") max_chars_number = gr.Number(label="Max Chars per Result", value=2000, precision=0, info="Maximum characters to retrieve per search result.") clear_button = gr.Button("Clear Chat") cancel_button = gr.Button("Cancel Generation") with gr.Column(scale=7): chatbot = gr.Chatbot(label="Chat", type="messages") msg_input = gr.Textbox(label="Your Message", placeholder="Enter your message and press Enter") search_debug = gr.Markdown(label="Web Search Debug") # Wire the Enable Web Search checkbox change to update the System Prompt textbox. enable_search_checkbox.change( fn=update_default_prompt, inputs=[enable_search_checkbox], outputs=[system_prompt_text] ) def clear_chat(): return [], "", "" clear_button.click(fn=clear_chat, outputs=[chatbot, msg_input, search_debug]) cancel_button.click(fn=cancel_generation, outputs=search_debug) # Submission: the chat_response function is used with streaming. msg_input.submit( fn=chat_response, inputs=[msg_input, chatbot, system_prompt_text, enable_search_checkbox, max_results_number, max_chars_number, model_dropdown, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repeat_penalty_slider], outputs=[chatbot, search_debug], ) demo.launch()