Create app.py

app.py

@@ -0,0 +1,448 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import spaces
+import re
+
+# Model configuration
+model_name = "HelpingAI/Dhanishtha-2.0-preview"
+
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
+
+def load_model():
+    """Load the model and tokenizer"""
+    global model, tokenizer
+
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Ensure pad token is set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True
+    )
+
+    print("Model loaded successfully!")
+
+def format_thinking_text(text):
+    """Format text to properly display <think> tags in Gradio with better styling"""
+    if not text:
+        return text
+
+    # More sophisticated formatting for thinking blocks
+    # Replace <think> and </think> tags with styled markdown
+    formatted_text = text
+
+    # Handle thinking blocks with proper markdown formatting
+    thinking_pattern = r'<think>(.*?)</think>'
+
+    def replace_thinking_block(match):
+        thinking_content = match.group(1).strip()
+        return f"\n\n💭 **Thinking Process:**\n\n```\n{thinking_content}\n```\n\n"
+
+    formatted_text = re.sub(thinking_pattern, replace_thinking_block, formatted_text, flags=re.DOTALL)
+
+    # Clean up any remaining raw tags that might not have been caught
+    formatted_text = re.sub(r'</?think>', '', formatted_text)
+
+    return formatted_text.strip()
+
+@spaces.GPU()
+def generate_response(message, history, max_tokens, temperature, top_p):
+    """Generate streaming response without threading"""
+    global model, tokenizer
+
+    if model is None or tokenizer is None:
+        yield "Model is still loading. Please wait..."
+        return
+
+    # Prepare conversation history
+    messages = []
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
+
+    # Add current message
+    messages.append({"role": "user", "content": message})
+
+    # Apply chat template
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    # Tokenize input
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    try:
+        with torch.no_grad():
+            # Use transformers streaming with custom approach
+            generated_text = ""
+            current_input_ids = model_inputs["input_ids"]
+            current_attention_mask = model_inputs["attention_mask"]
+
+            for _ in range(max_tokens):
+                # Generate next token
+                outputs = model(
+                    input_ids=current_input_ids,
+                    attention_mask=current_attention_mask,
+                    use_cache=True
+                )
+
+                # Get logits for the last token
+                logits = outputs.logits[0, -1, :]
+
+                # Apply temperature
+                if temperature != 1.0:
+                    logits = logits / temperature
+
+                # Apply top-p sampling
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
+                    sorted_indices_to_remove[0] = 0
+                    indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                    logits[indices_to_remove] = float('-inf')
+
+                # Sample next token
+                probs = torch.softmax(logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+
+                # Check for EOS token
+                if next_token.item() == tokenizer.eos_token_id:
+                    break
+
+                # Decode the new token (preserve special tokens like <think>)
+                new_token_text = tokenizer.decode(next_token, skip_special_tokens=False)
+                generated_text += new_token_text
+
+                # Format and yield the current text
+                formatted_text = format_thinking_text(generated_text)
+                yield formatted_text
+
+                # Update inputs for next iteration
+                current_input_ids = torch.cat([current_input_ids, next_token.unsqueeze(0)], dim=-1)
+                current_attention_mask = torch.cat([current_attention_mask, torch.ones((1, 1), device=model.device)], dim=-1)
+
+    except Exception as e:
+        yield f"Error generating response: {str(e)}"
+        return
+
+    # Final yield with complete formatted text
+    final_text = format_thinking_text(generated_text) if generated_text else "No response generated."
+    yield final_text
+
+def chat_interface(message, history, max_tokens, temperature, top_p):
+    """Main chat interface with improved streaming"""
+    if not message.strip():
+        return history, ""
+
+    # Add user message to history
+    history.append([message, ""])
+
+    # Generate response with streaming
+    for partial_response in generate_response(message, history[:-1], max_tokens, temperature, top_p):
+        history[-1][1] = partial_response
+        yield history, ""
+
+    return history, ""
+
+# Load model on startup
+print("Initializing model...")
+load_model()
+
+# Custom CSS for better styling and thinking blocks
+custom_css = """
+/* Main chatbot styling */
+.chatbot {
+    font-size: 14px;
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+}
+
+/* Thinking block styling */
+.thinking-block {
+    background: linear-gradient(135deg, #f0f8ff 0%, #e6f3ff 100%);
+    border-left: 4px solid #4a90e2;
+    border-radius: 8px;
+    padding: 12px 16px;
+    margin: 12px 0;
+    font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    position: relative;
+}
+
+.thinking-block::before {
+    content: "🤔";
+    position: absolute;
+    top: -8px;
+    left: 12px;
+    background: white;
+    padding: 0 4px;
+    font-size: 16px;
+}
+
+/* Message styling */
+.message {
+    padding: 10px 14px;
+    margin: 6px 0;
+    border-radius: 12px;
+    line-height: 1.5;
+}
+
+.user-message {
+    background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
+    margin-left: 15%;
+    border-bottom-right-radius: 4px;
+}
+
+.assistant-message {
+    background: linear-gradient(135deg, #f5f5f5 0%, #eeeeee 100%);
+    margin-right: 15%;
+    border-bottom-left-radius: 4px;
+}
+
+/* Code block styling */
+pre {
+    background-color: #f8f9fa;
+    border: 1px solid #e9ecef;
+    border-radius: 6px;
+    padding: 12px;
+    overflow-x: auto;
+    font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
+    font-size: 13px;
+    line-height: 1.4;
+}
+
+/* Button styling */
+.gradio-button {
+    border-radius: 8px;
+    font-weight: 500;
+    transition: all 0.2s ease;
+}
+
+.gradio-button:hover {
+    transform: translateY(-1px);
+    box-shadow: 0 4px 8px rgba(0,0,0,0.15);
+}
+
+/* Input styling */
+.gradio-textbox {
+    border-radius: 8px;
+    border: 2px solid #e0e0e0;
+    transition: border-color 0.2s ease;
+}
+
+.gradio-textbox:focus {
+    border-color: #4a90e2;
+    box-shadow: 0 0 0 3px rgba(74, 144, 226, 0.1);
+}
+
+/* Slider styling */
+.gradio-slider {
+    margin: 8px 0;
+}
+
+/* Examples styling */
+.gradio-examples {
+    margin-top: 16px;
+}
+
+.gradio-examples .gradio-button {
+    background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%);
+    border: 1px solid #dee2e6;
+    color: #495057;
+    font-size: 13px;
+    padding: 8px 12px;
+}
+
+.gradio-examples .gradio-button:hover {
+    background: linear-gradient(135deg, #e9ecef 0%, #dee2e6 100%);
+    color: #212529;
+}
+"""
+
+# Create Gradio interface
+with gr.Blocks(
+    title="🤖 Dhanishtha-2.0-preview Chat",
+    theme=gr.themes.Soft(),
+    css=custom_css
+) as demo:
+    gr.Markdown(
+        """
+        # 🤖 Dhanishtha-2.0-preview Chat
+
+        Chat with the **HelpingAI/Dhanishtha-2.0-preview** model - The world's first LLM designed to think between responses!
+
+        ### ✨ Key Features:
+        - 🧠 **Multi-step Reasoning**: Unlike other LLMs that think once, Dhanishtha can think, rethink, self-evaluate, and refine using multiple `<think>` blocks
+        - 🔄 **Iterative Thinking**: Watch the model's thought process unfold in real-time
+        - 💡 **Enhanced Problem Solving**: Better reasoning capabilities through structured thinking
+
+        **Note**: The `<think>` blocks show the model's internal reasoning process and will be displayed in a formatted way below.
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(
+                [],
+                elem_id="chatbot",
+                bubble_full_width=False,
+                height=600,
+                show_copy_button=True,
+                show_share_button=True,
+                avatar_images=("👤", "🤖"),
+                render_markdown=True,
+                latex_delimiters=[
+                    {"left": "$$", "right": "$$", "display": True},
+                    {"left": "$", "right": "$", "display": False}
+                ]
+            )
+
+            with gr.Row():
+                msg = gr.Textbox(
+                    container=False,
+                    placeholder="Ask me anything! I'll show you my thinking process...",
+                    label="Message",
+                    autofocus=True,
+                    scale=8,
+                    lines=1,
+                    max_lines=5
+                )
+                send_btn = gr.Button("🚀 Send", variant="primary", scale=1, size="lg")
+
+        with gr.Column(scale=1, min_width=300):
+            gr.Markdown("### ⚙️ Generation Parameters")
+
+            max_tokens = gr.Slider(
+                minimum=50,
+                maximum=8192,
+                value=2048,
+                step=50,
+                label="🎯 Max Tokens",
+                info="Maximum number of tokens to generate"
+            )
+
+            temperature = gr.Slider(
+                minimum=0.1,
+                maximum=2.0,
+                value=0.7,
+                step=0.1,
+                label="🌡️ Temperature",
+                info="Higher = more creative, Lower = more focused"
+            )
+
+            top_p = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.9,
+                step=0.05,
+                label="🎲 Top-p",
+                info="Nucleus sampling threshold"
+            )
+
+            with gr.Row():
+                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary", scale=1)
+                stop_btn = gr.Button("⏹️ Stop", variant="stop", scale=1)
+
+            gr.Markdown("### 📊 Model Info")
+            gr.Markdown(
+                """
+                **Model**: HelpingAI/Dhanishtha-2.0-preview
+                **Type**: Reasoning LLM with thinking blocks
+                **Features**: Multi-step reasoning, self-evaluation
+                """
+            )
+    
+    # Event handlers
+    def submit_message(message, history, max_tokens, temperature, top_p):
+        """Handle message submission"""
+        return chat_interface(message, history, max_tokens, temperature, top_p)
+
+    def clear_chat():
+        """Clear the chat history"""
+        return [], ""
+
+    # Message submission events
+    msg.submit(
+        submit_message,
+        inputs=[msg, chatbot, max_tokens, temperature, top_p],
+        outputs=[chatbot, msg],
+        concurrency_limit=1,
+        show_progress="minimal"
+    )
+
+    send_btn.click(
+        submit_message,
+        inputs=[msg, chatbot, max_tokens, temperature, top_p],
+        outputs=[chatbot, msg],
+        concurrency_limit=1,
+        show_progress="minimal"
+    )
+
+    # Clear chat event
+    clear_btn.click(
+        clear_chat,
+        outputs=[chatbot, msg],
+        show_progress=False
+    )
+
+    # Example prompts section
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                ["Hello! Can you introduce yourself and show me how you think?"],
+                ["Solve this step by step: What is 15% of 240?"],
+                ["Explain quantum entanglement in simple terms"],
+                ["Write a short Python function to find the factorial of a number"],
+                ["What are the pros and cons of renewable energy?"],
+                ["Help me understand the difference between AI and machine learning"],
+                ["Create a haiku about artificial intelligence"],
+                ["Explain why the sky is blue using physics principles"]
+            ],
+            inputs=msg,
+            label="💡 Example Prompts - Try these to see the thinking process!",
+            examples_per_page=4
+        )
+
+    # Footer with information
+    gr.Markdown(
+        """
+        ---
+        ### 🔧 Technical Details
+        - **Model**: HelpingAI/Dhanishtha-2.0-preview
+        - **Framework**: Transformers + Gradio
+        - **Features**: Real-time streaming, thinking process visualization, custom sampling
+        - **Reasoning**: Multi-step thinking with `<think>` blocks for transparent AI reasoning
+
+        **Note**: This interface streams responses token by token and formats thinking blocks for better readability.
+        The model's internal reasoning process is displayed in formatted code blocks.
+
+        ---
+        *Built with ❤️ using Gradio and Transformers*
+        """
+    )
+
+if __name__ == "__main__":
+    demo.queue(
+        max_size=20,
+        default_concurrency_limit=1
+    ).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        quiet=False
+    )