app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import spaces

class VibeThinker:
    def __init__(self, model_path="WeiboAI/VibeThinker-1.5B"):
        self.model_path = model_path
        print("Loading model... This may take a minute.")
        
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_path,
            low_cpu_mem_usage=True,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            trust_remote_code=True
        )
        
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_path, 
            trust_remote_code=True
        )
        
        print(f"Model loaded successfully!")
        print(f"Using device: {self.model.device}")
        if torch.cuda.is_available():
            print(f"CUDA device: {torch.cuda.get_device_name(0)}")

    @spaces.GPU  # This decorator allocates GPU when function is called (for ZeroGPU spaces)
    def infer_text(self, prompt, temperature=0.6, max_tokens=40960, top_p=0.95):
        """
        Generate response for a given prompt
        
        Args:
            prompt: The input question (preferably in English)
            temperature: Controls randomness (0.6 or 1.0 recommended)
            max_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter
        """
        messages = [
            {"role": "user", "content": prompt}
        ]
        
        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
        
        generation_config = dict(
            max_new_tokens=max_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            top_k=None  # Set to -1 in vLLM/SGLang
        )
        
        print(f"Generating response with temperature={temperature}, max_tokens={max_tokens}...")
        
        generated_ids = self.model.generate(
            **model_inputs,
            generation_config=GenerationConfig(**generation_config)
        )
        
        generated_ids = [
            output_ids[len(input_ids):] 
            for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]
        
        response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        return response


# Initialize model
print("Initializing VibeThinker-1.5B...")
model = VibeThinker()

# Create Gradio interface
def generate_response(prompt, temperature, max_tokens, top_p):
    if not prompt.strip():
        return "Please enter a question."
    
    try:
        response = model.infer_text(
            prompt=prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p
        )
        return response
    except Exception as e:
        return f"Error: {str(e)}"


# Gradio UI
with gr.Blocks(title="VibeThinker-1.5B Math & Code Reasoning") as demo:
    gr.Markdown("""
    # 🧠 VibeThinker-1.5B: Reasoning Model
    
    **Optimized for**: Competitive math problems and algorithm coding challenges
    
    **Note**: This model works best with questions in English. It's specifically trained for 
    mathematical reasoning and competitive programming tasks.
    
    ### Example Prompts:
    - "Solve: Find all solutions to x^3 - 3x^2 + 4 = 0"
    - "Write a Python function to find the longest palindromic substring in O(n^2) time"
    - "Prove that the sum of angles in a triangle equals 180 degrees"
    
    [GitHub](https://github.com/WeiboAI/VibeThinker) | [HuggingFace Model](https://huggingface.co/WeiboAI/VibeThinker-1.5B) | [Paper](https://huggingface.co/papers/2511.06221)
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Your Question",
                placeholder="Enter your math problem or coding challenge here (in English)...",
                lines=5
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                temperature_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.5,
                    value=0.6,
                    step=0.1,
                    label="Temperature (0.6 or 1.0 recommended)"
                )
                
                max_tokens_slider = gr.Slider(
                    minimum=512,
                    maximum=40960,
                    value=8192,
                    step=512,
                    label="Max Tokens (model supports up to 40,960)"
                )
                
                top_p_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top P"
                )
            
            submit_btn = gr.Button("🚀 Generate Solution", variant="primary")
            clear_btn = gr.Button("🗑️ Clear")
        
        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="Model Response",
                lines=20,
                show_copy_button=True
            )
    
    # Example questions
    gr.Examples(
        examples=[
            ["Find the number of positive integers n ≤ 1000 such that n^2 + n + 41 is prime.", 0.6, 8192, 0.95],
            ["Write an efficient algorithm to solve the 0-1 knapsack problem using dynamic programming.", 0.6, 8192, 0.95],
            ["Prove that √2 is irrational using proof by contradiction.", 0.6, 8192, 0.95],
            ["A tank can be filled by pipe A in 3 hours and pipe B in 5 hours. If both pipes are opened together, how long will it take to fill the tank?", 0.6, 8192, 0.95],
        ],
        inputs=[prompt_input, temperature_slider, max_tokens_slider, top_p_slider],
        label="Example Problems"
    )
    
    # Event handlers
    submit_btn.click(
        fn=generate_response,
        inputs=[prompt_input, temperature_slider, max_tokens_slider, top_p_slider],
        outputs=output_text
    )
    
    clear_btn.click(
        fn=lambda: ("", ""),
        inputs=[],
        outputs=[prompt_input, output_text]
    )
    
    gr.Markdown("""
    ---
    ### 📊 Model Performance Highlights:
    - **AIME24**: 80.3 (vs DeepSeek R1: 79.8)
    - **AIME25**: 74.4 (vs DeepSeek R1: 70.0)  
    - **LiveCodeBench v6**: 51.1
    - **Parameters**: Only 1.5B (400x smaller than DeepSeek R1!)
    
    **Training Cost**: $7,800 USD | **License**: MIT
    """)

# Launch the app
if __name__ == "__main__":
    demo.queue()  # Enable queuing for better UX
    demo.launch()