Spaces:
Running
Running
| import os | |
| # Fixes the Gradio Analytics crash bug on Colab/Spaces | |
| os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" | |
| import torch | |
| import gc | |
| import re | |
| import threading | |
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| from peft import PeftModel | |
| # ========================================== | |
| # 1. SMART PRE-LOAD MODELS (NO QUANTIZATION) | |
| # ========================================== | |
| if "loaded_engines" not in globals(): | |
| global loaded_engines | |
| loaded_engines = {} | |
| MODELS_CONFIG = { | |
| "ReasonBorn-Instruct": { | |
| "base": "Qwen/Qwen2.5-3B-Instruct", | |
| "adapter": "Phase-Technologies/ReasonBorn-Qwen-3B", | |
| }, | |
| "ReasonBorn-LoRA": { | |
| "base": "Qwen/Qwen2.5-3B", | |
| "adapter": "Phase-Technologies/rb-qwen3b-16ds-lora", | |
| } | |
| } | |
| if not loaded_engines: | |
| print("Initializing Xerv Systems... Pre-loading models for instant streaming.") | |
| # Force single-device mapping to prevent PEFT offload KeyError | |
| target_device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Targeting inference device: {target_device.upper()}") | |
| for key, cfg in MODELS_CONFIG.items(): | |
| print(f"--- Loading {key} (Unquantized BF16) ---") | |
| tokenizer = AutoTokenizer.from_pretrained(cfg["adapter"]) | |
| # Load Base Model on a single device to avoid meta-tensor offloading issues | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| cfg["base"], | |
| torch_dtype=torch.bfloat16, | |
| device_map={"": target_device}, | |
| trust_remote_code=True | |
| ) | |
| # Merge adapter for inference | |
| model = PeftModel.from_pretrained(base_model, cfg["adapter"]) | |
| model.eval() | |
| loaded_engines[key] = {"model": model, "tokenizer": tokenizer} | |
| print("✅ Both Reasoning Engines successfully loaded and ready.") | |
| else: | |
| print("⚡ Models already detected in memory! Skipping load phase for instant boot.") | |
| # ========================================== | |
| # 2. BULLETPROOF LATEX & TAG PARSER | |
| # ========================================== | |
| def format_output_with_latex_support(text): | |
| # Standardize LaTeX delimiters for Gradio | |
| text = text.replace(r'\\(', '$').replace(r'\\)', '$') | |
| text = text.replace(r'\\[', '$$').replace(r'\\]', '$$') | |
| # Extract Conclusion | |
| conclusion_match = re.search(r"<conclusion>(.*?)(?:</conclusion>|$)", text, re.DOTALL) | |
| if conclusion_match: | |
| conclusion_text = conclusion_match.group(1).strip() | |
| thinking_text = text[:conclusion_match.start()].strip() | |
| # Format Thinking Process | |
| thinking_text = thinking_text.replace("<plan>", "**🔹 PLAN:**\n").replace("</plan>", "\n") | |
| thinking_text = thinking_text.replace("<reasoning>", "\n").replace("</reasoning>", "\n") | |
| # Handle dynamic <step> tags | |
| thinking_text = re.sub(r"<step(?:\s+index=\"(\d+)\")?>", | |
| lambda m: f"**🔸 STEP {m.group(1)}:** " if m.group(1) else "**🔸 STEP:** ", | |
| thinking_text) | |
| thinking_text = thinking_text.replace("</step>", "\n") | |
| thinking_text = thinking_text.replace("<verify>", "**✅ VERIFY:** ").replace("</verify>", "\n") | |
| # Wrap thinking in a collapsible HTML details block | |
| formatted = ( | |
| f"<details>\n" | |
| f"<summary>🧠 View Thinking Process</summary>\n\n" | |
| f"{thinking_text}\n\n" | |
| f"</details>\n\n" | |
| f"**🎯 CONCLUSION:**\n\n{conclusion_text}" | |
| ) | |
| return formatted | |
| else: | |
| # Fallback if generation stops before conclusion | |
| text = text.replace("<plan>", "**🔹 PLAN:**\n").replace("</plan>", "\n") | |
| text = text.replace("<reasoning>", "\n").replace("</reasoning>", "\n") | |
| text = re.sub(r"<step(?:\s+index=\"(\d+)\")?>", | |
| lambda m: f"**🔸 STEP {m.group(1)}:** " if m.group(1) else "**🔸 STEP:** ", | |
| text) | |
| text = text.replace("</step>", "\n") | |
| text = text.replace("<verify>", "**✅ VERIFY:** ").replace("</verify>", "\n") | |
| return text | |
| # ========================================== | |
| # 3. REAL-TIME STREAMING GENERATOR | |
| # ========================================== | |
| def process_chat_stream(user_message, history, model_choice): | |
| """ | |
| Handles Gradio's 'messages' format natively: [{"role": "user", "content": "..."}, ...] | |
| """ | |
| if not user_message.strip(): | |
| yield "", gr.update(), gr.update(), gr.update() | |
| return | |
| # Initialize history if empty and append new user/assistant dicts | |
| history = history or [] | |
| history.append({"role": "user", "content": user_message}) | |
| history.append({"role": "assistant", "content": ""}) | |
| # Yield immediately to update UI (hide hero/suggestions, show chatbot) | |
| yield "", gr.update(value=history, visible=True), gr.update(visible=False), gr.update(visible=False) | |
| try: | |
| engine = loaded_engines[model_choice] | |
| model = engine["model"] | |
| tokenizer = engine["tokenizer"] | |
| # Build strict ReasonBorn System Prompt | |
| prompt = "<|im_start|>system\nYou are ReasonBorn. Use <plan>, <reasoning> with <step> & <verify>, <conclusion> strictly.<|im_end|>\n" | |
| # Append prior conversation history (excluding the two entries we just appended) | |
| for msg in history[:-2]: | |
| role = msg["role"] | |
| content = msg["content"] | |
| if role == "user": | |
| prompt += f"<|im_start|>user\n{content}<|im_end|>\n" | |
| elif role == "assistant": | |
| # Strip out HTML UI elements so the model only sees plain text history | |
| clean_content = re.sub(r"<.*?>", "", content) | |
| prompt += f"<|im_start|>assistant\n{clean_content}<|im_end|>\n" | |
| # Append current message | |
| prompt += f"<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant\n" | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) | |
| generation_kwargs = dict( | |
| **inputs, | |
| max_new_tokens=1024, | |
| temperature=0.2, | |
| top_p=0.9, | |
| repetition_penalty=1.1, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id, | |
| eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>"), | |
| streamer=streamer | |
| ) | |
| # Start generation in a separate thread | |
| thread = threading.Thread(target=model.generate, kwargs=generation_kwargs) | |
| thread.start() | |
| accumulated_text = "" | |
| # Stream chunks back to UI | |
| for new_text in streamer: | |
| accumulated_text += new_text | |
| # Real-time formatting for visual feedback | |
| live_text = accumulated_text.replace(r'\\(', '$').replace(r'\\)', '$').replace(r'\\[', '$$').replace(r'\\]', '$$') | |
| live_text = live_text.replace("<plan>", "**🔹 PLAN:**\n").replace("</plan>", "\n") | |
| live_text = live_text.replace("<reasoning>", "\n").replace("</reasoning>", "\n") | |
| live_text = re.sub(r"<step(?:\s+index=\"(\d+)\")?>", | |
| lambda m: f"**🔸 STEP {m.group(1)}:** " if m.group(1) else "**🔸 STEP:** ", | |
| live_text) | |
| live_text = live_text.replace("</step>", "\n") | |
| live_text = live_text.replace("<verify>", "**✅ VERIFY:** ").replace("</verify>", "\n") | |
| live_text = live_text.replace("<conclusion>", "\n\n**🎯 CONCLUSION:**\n\n").replace("</conclusion>", "") | |
| # Update the latest bot message in history dictionaries | |
| history[-1]["content"] = live_text + " ⏳" | |
| yield "", gr.update(value=history, visible=True), gr.update(visible=False), gr.update(visible=False) | |
| # Final formatting pass with HTML block wrapping | |
| final_formatted = format_output_with_latex_support(accumulated_text) | |
| history[-1]["content"] = final_formatted | |
| yield "", gr.update(value=history, visible=True), gr.update(visible=False), gr.update(visible=False) | |
| # Cleanup memory | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| except Exception as e: | |
| history[-1]["content"] = f"**System Error:** {str(e)}" | |
| yield "", gr.update(value=history, visible=True), gr.update(visible=False), gr.update(visible=False) | |
| # ========================================== | |
| # 4. UI/UX: ADAPTIVE DARK/LIGHT MODE CSS | |
| # ========================================== | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Google+Sans:wght@400;500;700&display=swap'); | |
| /* Global Typography & Layout */ | |
| .gradio-container { font-family: 'Google Sans', sans-serif !important; } | |
| .main-wrap { max-width: 750px !important; margin: 0 auto !important; padding-bottom: 100px !important; } | |
| /* Hero Section */ | |
| .xerv-title { font-size: 46px; font-weight: 700; letter-spacing: -1px; margin-top: 40px; margin-bottom: 8px;} | |
| .greeting { font-size: 18px; margin-bottom: 4px; opacity: 0.7;} | |
| .subtitle { font-size: 26px; font-weight: 500; margin-bottom: 30px;} | |
| /* Chat Window Base */ | |
| #chat-window { height: 65vh !important; } | |
| /* User Bubble - Always Blue */ | |
| .message.user { background: #2563eb !important; color: white !important; border-radius: 20px 20px 0 20px !important; padding: 14px 20px !important; font-size: 16px !important; } | |
| .message.user * { color: white !important; } | |
| /* Bot Bubble - Light Mode (Default) */ | |
| .message.bot { background: #ffffff !important; color: #0f172a !important; border: 1px solid #e2e8f0 !important; border-radius: 20px 20px 20px 0 !important; padding: 16px 20px !important; font-size: 16px !important; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05) !important; } | |
| /* Bot Bubble - Dark Mode */ | |
| .dark .message.bot { background: #1e293b !important; color: #f8fafc !important; border-color: #334155 !important; } | |
| /* Thinking Details Block - Light Mode */ | |
| #chat-window details { background-color: #f8fafc !important; border: 1px solid #e2e8f0 !important; border-radius: 12px !important; padding: 14px !important; margin-bottom: 16px !important; box-shadow: inset 0 2px 4px 0 rgb(0 0 0 / 0.02) !important; transition: all 0.2s ease !important; } | |
| #chat-window summary { cursor: pointer !important; font-weight: 600 !important; font-size: 15px !important; user-select: none !important; outline: none !important; color: #334155 !important;} | |
| /* Thinking Details Block - Dark Mode */ | |
| .dark #chat-window details { background-color: #0f172a !important; border-color: #1e293b !important; color: #cbd5e1 !important; } | |
| .dark #chat-window summary { color: #94a3b8 !important; } | |
| #chat-window details[open] summary { margin-bottom: 12px !important; padding-bottom: 12px !important; border-bottom: 1px solid rgba(128,128,128,0.2) !important; } | |
| /* Input Row - Adaptive */ | |
| .input-row { align-items: center !important; border-radius: 30px !important; padding: 6px 14px !important; border: 1px solid #cbd5e1 !important; transition: all 0.2s; box-shadow: 0 4px 6px -1px rgba(0,0,0,0.05) !important; background: #f8fafc !important; } | |
| .dark .input-row { background: #1e293b !important; border-color: #334155 !important; } | |
| .input-row:focus-within { border-color: #3b82f6 !important; box-shadow: 0 4px 12px rgba(59, 130, 246, 0.15) !important; } | |
| .input-row textarea { background: transparent !important; border: none !important; box-shadow: none !important; font-size: 16px !important; } | |
| .input-row textarea:focus { outline: none !important; border: none !important; box-shadow: none !important; } | |
| /* Buttons */ | |
| .send-button { background: #2563eb !important; color: white !important; border-radius: 50% !important; height: 42px !important; width: 42px !important; min-width: 42px !important; padding: 0 !important; border: none !important; display: flex; justify-content: center; align-items: center; } | |
| .send-button:disabled { background: #94a3b8 !important; } | |
| .dark .send-button:disabled { background: #334155 !important; color: #64748b !important; } | |
| /* Suggestions - Adaptive */ | |
| .sugg-btn { background: #ffffff !important; border: 1px solid #e2e8f0 !important; border-radius: 16px !important; padding: 16px 20px !important; text-align: left !important; justify-content: flex-start !important; font-size: 16px !important; color: #1e293b !important; box-shadow: 0 1px 2px rgba(0,0,0,0.05) !important; margin-bottom: 12px !important; cursor: pointer !important; } | |
| .dark .sugg-btn { background: #1e293b !important; border-color: #334155 !important; color: #f8fafc !important; } | |
| .sugg-btn:hover { opacity: 0.8; } | |
| /* LaTeX Fixes */ | |
| .katex-display { margin: 1em 0 !important; overflow-x: auto !important; overflow-y: hidden !important; padding: 8px 0 !important; } | |
| .katex { font-size: 1.1em !important; } | |
| footer, .label-wrap { display: none !important; } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_classes="main-wrap"): | |
| with gr.Column(elem_id="hero-section") as hero: | |
| gr.HTML(""" | |
| <div class="xerv-title">Xerv</div> | |
| <div class="greeting">Hey there!</div> | |
| <div class="subtitle">Let's make something happen.</div> | |
| """) | |
| with gr.Column(elem_id="suggestions-section") as suggestions: | |
| btn1 = gr.Button(r"🔍 Prove that $\sqrt{2}$ is irrational", elem_classes="sugg-btn") | |
| btn2 = gr.Button(r"🧮 Solve $x^3 - 6x^2 + 11x - 6 = 0$", elem_classes="sugg-btn") | |
| btn3 = gr.Button(r"📊 Explain eigenvalues with a matrix example", elem_classes="sugg-btn") | |
| chatbot = gr.Chatbot( | |
| visible=False, | |
| elem_id="chat-window", | |
| show_label=False, | |
| avatar_images=(None, None), | |
| sanitize_html=False, | |
| # Note: Removed type="messages" to resolve the TypeError in Gradio 6.0 | |
| latex_delimiters=[ | |
| {"left": "$$", "right": "$$", "display": True}, | |
| {"left": "$", "right": "$", "display": False} | |
| ] | |
| ) | |
| with gr.Column(): | |
| with gr.Row(elem_classes="input-row"): | |
| chat_input = gr.Textbox( | |
| show_label=False, | |
| placeholder="Ask Xerv to solve complex math...", | |
| lines=1, | |
| max_lines=4, | |
| scale=8 | |
| ) | |
| send_btn = gr.Button("🚀", elem_classes="send-button", scale=1) | |
| model_selector = gr.Radio( | |
| choices=list(MODELS_CONFIG.keys()), | |
| value="ReasonBorn-Instruct", | |
| label="Reasoning Engine", | |
| container=False | |
| ) | |
| # --- Wire up Interactivity --- | |
| chat_input.submit( | |
| process_chat_stream, | |
| inputs=[chat_input, chatbot, model_selector], | |
| outputs=[chat_input, chatbot, hero, suggestions] | |
| ) | |
| send_btn.click( | |
| process_chat_stream, | |
| inputs=[chat_input, chatbot, model_selector], | |
| outputs=[chat_input, chatbot, hero, suggestions] | |
| ) | |
| btn1.click( | |
| fn=lambda: r"Prove that $\sqrt{2}$ is irrational using step-by-step logic", | |
| outputs=[chat_input] | |
| ).then( | |
| fn=process_chat_stream, | |
| inputs=[chat_input, chatbot, model_selector], | |
| outputs=[chat_input, chatbot, hero, suggestions] | |
| ) | |
| btn2.click( | |
| fn=lambda: r"Solve $x^3 - 6x^2 + 11x - 6 = 0$ and verify roots", | |
| outputs=[chat_input] | |
| ).then( | |
| fn=process_chat_stream, | |
| inputs=[chat_input, chatbot, model_selector], | |
| outputs=[chat_input, chatbot, hero, suggestions] | |
| ) | |
| btn3.click( | |
| fn=lambda: r"Explain eigenvalues in linear algebra with an example matrix", | |
| outputs=[chat_input] | |
| ).then( | |
| fn=process_chat_stream, | |
| inputs=[chat_input, chatbot, model_selector], | |
| outputs=[chat_input, chatbot, hero, suggestions] | |
| ) | |
| if __name__ == "__main__": | |
| # Removed the manual light mode javascript. Added adaptive CSS directly to launch parameters. | |
| demo.launch( | |
| share=True, | |
| debug=True, | |
| css=CSS, | |
| theme=gr.themes.Default() | |
| ) | |