import subprocess subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True) import torch import spaces import gradio as gr import flash_attn from threading import Thread from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList ) MODEL_ID = "unsloth/QwQ-32B-unsloth-bnb-4bit" DEFAULT_SYSTEM_PROMPT = """ Think step by step and explain your reasoning clearly. Break down the problem into logical components, verify each step, and ensure consistency before arriving at the final answer." For complex reasoning tasks, you can enhance it with: "If there are multiple possible solutions, consider each one before selecting the best answer." "Use intermediate calculations and justify each step before proceeding." "If relevant, include real-world analogies to improve clarity. """ CSS = """ .gr-chatbot { min-height: 500px; border-radius: 15px; } .special-tag { color: #2ecc71; font-weight: 600; } footer { display: none !important; } """ class StopOnTokens(StoppingCriteria): def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: # Stop when the EOS token is generated. return input_ids[0][-1] == tokenizer.eos_token_id def initialize_model(): # Enable 4-bit quantization for faster inference and lower memory usage. quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, ) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="cuda", #quantization_config=quantization_config, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="flash_attention_2" ) model.to("cuda") model.eval() # set evaluation mode to disable gradients and speed up inference return model, tokenizer def format_response(text): # List of replacements to format key tokens with HTML for styling. replacements = [ ("[Understand]", '\n[Understand]\n'), ("[think]", '\n[think]\n'), ("[/think]", '\n[/think]\n'), ("[Answer]", '\n[Answer]\n'), ("[/Answer]", '\n[/Answer]\n'), ] for old, new in replacements: text = text.replace(old, new) return text # --- New helper: Llama-3 conversation template --- def apply_llama3_chat_template(conversation, add_generation_prompt=True): """ Convert the conversation (a list of dicts with 'role' and 'content') into a single prompt string in Llama-3 style. """ prompt = "" for msg in conversation: role = msg["role"].upper() if role == "SYSTEM": prompt += "<|SYSTEM|>\n" + msg["content"].strip() + "\n" elif role == "USER": prompt += "<|USER|>\n" + msg["content"].strip() + "\n" elif role == "ASSISTANT": prompt += "<|ASSISTANT|>\n" + msg["content"].strip() + "\n" if add_generation_prompt: prompt += "<|ASSISTANT|>\n" return prompt @spaces.GPU(duration=120) def generate_response(message, chat_history, system_prompt, temperature, max_tokens, top_p, top_k, repetition_penalty): # Build the conversation history. conversation = [{"role": "system", "content": system_prompt}] for user_msg, bot_msg in chat_history: conversation.append({"role": "user", "content": user_msg}) conversation.append({"role": "assistant", "content": bot_msg}) conversation.append({"role": "user", "content": message}) # Use the Llama-3 conversation template to build the prompt. prompt = apply_llama3_chat_template(conversation, add_generation_prompt=True) input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device) # Setup the streamer to yield new tokens as they are generated. streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True) # Prepare generation parameters including extra customization options. generate_kwargs = { "input_ids": input_ids, "streamer": streamer, "max_new_tokens": max_tokens, "temperature": temperature, "top_p": top_p, "top_k": top_k, "repetition_penalty": repetition_penalty, "stopping_criteria": StoppingCriteriaList([StopOnTokens()]) } # Run the generation inside a no_grad block for speed. def generate_inference(): with torch.inference_mode(): model.generate(**generate_kwargs) Thread(target=generate_inference, daemon=True).start() # Stream the output tokens. partial_message = "" new_history = chat_history + [(message, "")] for new_token in streamer: partial_message += new_token formatted = format_response(partial_message) new_history[-1] = (message, formatted + "▌") yield new_history # Final update without the cursor. new_history[-1] = (message, format_response(partial_message)) yield new_history # Initialize the model and tokenizer globally. model, tokenizer = initialize_model() with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo: gr.Markdown("""

🧠 AI Reasoning Assistant

Ask me hard questions and see the reasoning unfold.

""") chatbot = gr.Chatbot(label="Conversation", elem_id="chatbot") msg = gr.Textbox(label="Your Question", placeholder="Type your question...") with gr.Accordion("⚙️ Settings", open=False): system_prompt = gr.TextArea(value=DEFAULT_SYSTEM_PROMPT, label="System Instructions") temperature = gr.Slider(0, 1, value=0.6, label="Creativity (Temperature)") max_tokens = gr.Slider(128, 32768, 32768, label="Max Response Length") top_p = gr.Slider(0.0, 1.0, value=0.95, label="Top P (Nucleus Sampling)") top_k = gr.Slider(0, 100, value=35, label="Top K") repetition_penalty = gr.Slider(0.5, 2.0, value=1.1, label="Repetition Penalty") clear = gr.Button("Clear History") # Link the input textbox with the generation function. msg.submit( generate_response, [msg, chatbot, system_prompt, temperature, max_tokens, top_p, top_k, repetition_penalty], chatbot, show_progress=True ) clear.click(lambda: None, None, chatbot, queue=False) if __name__ == "__main__": demo.queue().launch()