Update app.py

app.py

@@ -1,145 +1,202 @@
-import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import time
 
-MODEL_ID = "FractalAIResearch/Fathom-R1-14B"
-
-# Load model and tokenizer OUTSIDE the GPU function (following official docs)
-print("Loading model and tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True
-)
+# Load model and tokenizer
+model_name = "FractalAIResearch/Fathom-R1-14B"
 
+# Initialize tokenizer (can be done outside GPU decorator)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-# Move to GPU (following official docs pattern)
-model.to('cuda')
-print("Model loaded and moved to GPU")
+# Global model variable
+model = None
 
-@spaces.GPU
-def generate_response(message, history, max_tokens, temperature):
-    try:
-        # Simple prompt format
-        prompt = f"User: {message}\nAssistant:"
-        
-        # Tokenize
-        inputs = tokenizer(prompt, return_tensors="pt")
-        inputs = {k: v.to('cuda') for k, v in inputs.items()}
+def load_model():
+    """Load model on GPU"""
+    global model
+    if model is None:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+    return model
+
+@spaces.GPU #(duration=120)  # Allow up to 2 minutes for generation
+def generate_response(message, history, max_tokens=1024, temperature=0.7, top_p=0.9):
+    """Generate response using Fathom-R1-14B"""
+    
+    # Load model on GPU
+    model = load_model()
+    
+    # Format conversation history
+    conversation = []
+    for exchange in history:
+        if exchange['role'] == 'user':
+            conversation.append(f"User: {exchange['content']}")
+        else:
+            conversation.append(f"Assistant: {exchange['content']}")
+    
+    # Add current message
+    conversation.append(f"User: {message}")
+    conversation.append("Assistant:")
+    
+    # Create prompt
+    prompt = "\n".join(conversation)
+    
+    # Tokenize
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate with streaming
+    with torch.no_grad():
+        streamer_output = ""
         
-        # Generate
-        with torch.no_grad():
+        # Generate tokens one by one for streaming effect
+        for _ in range(max_tokens):
             outputs = model.generate(
-                **inputs,
-                max_new_tokens=max_tokens,
+                inputs,
+                max_new_tokens=1,
                 temperature=temperature,
+                top_p=top_p,
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
             )
-        
-        # Decode response
-        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
-        
-        # Update history
-        history.append([message, response])
-        return history, history, ""
-        
-    except Exception as e:
-        error_msg = f"Error: {str(e)}"
-        history.append([message, error_msg])
-        return history, history, ""
-
-# Create Gradio interface
-with gr.Blocks(title="Fathom R1 14B Chatbot") as demo:
-    gr.HTML("<h1>🤖 Fathom R1 14B Chatbot</h1>")
-    
-    with gr.Row():
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(height=500, label="Conversation")
             
-            with gr.Row():
-                msg = gr.Textbox(
-                    placeholder="Type your message here...", 
-                    label="Message",
-                    lines=3,
-                    scale=4
-                )
-                send_btn = gr.Button("Send", variant="primary", scale=1)
+            # Get new token
+            new_token = outputs[0, -1:]
+            new_text = tokenizer.decode(new_token, skip_special_tokens=True)
+            
+            # Check for end of sequence
+            if new_token.item() == tokenizer.eos_token_id:
+                break
+                
+            streamer_output += new_text
+            inputs = outputs
             
-            clear_btn = gr.Button("Clear Chat")
+            # Yield partial response for streaming
+            yield streamer_output
             
-        with gr.Column(scale=1):
-            gr.Markdown("### Settings")
+            # Small delay for streaming effect
+            time.sleep(0.05)
+
+# Alternative non-streaming version for faster response
+@spaces.GPU(duration=60)
+def generate_response_fast(message, history, max_tokens=1024, temperature=0.7, top_p=0.9):
+    """Generate response quickly without streaming"""
+    
+    # Load model on GPU
+    model = load_model()
+    
+    # Format conversation history
+    conversation = []
+    for exchange in history:
+        if exchange['role'] == 'user':
+            conversation.append(f"User: {exchange['content']}")
+        else:
+            conversation.append(f"Assistant: {exchange['content']}")
+    
+    # Add current message
+    conversation.append(f"User: {message}")
+    conversation.append("Assistant:")
+    
+    # Create prompt
+    prompt = "\n".join(conversation)
+    
+    # Tokenize
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode response
+    response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+    return response.strip()
+
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="Fathom-R1-14B Chatbot") as demo:
+        gr.Markdown("""
+        # 🧠 Fathom-R1-14B Reasoning Chatbot
+        
+        Powered by **FractalAI Research's Fathom-R1-14B** - a 14B parameter model optimized for mathematical and scientific reasoning tasks.
+        
+        This model excels at:
+        - Complex mathematical problems
+        - Scientific reasoning
+        - Step-by-step problem solving
+        - Logical analysis
+        """)
+        
+        with gr.Tab("Streaming Chat"):
+            chat_streaming = gr.ChatInterface(
+                fn=generate_response,
+                type="messages",
+                title="Streaming Response",
+                description="Get streaming responses (slower but shows progress)",
+                examples=[
+                    "Solve this math problem: If a train travels 120 km in 2 hours, what's its average speed?",
+                    "Explain the concept of photosynthesis step by step",
+                    "What is the derivative of x^3 + 2x^2 - 5x + 3?",
+                    "How do you calculate the area of a circle with radius 7?"
+                ]
+            )
+        
+        with gr.Tab("Fast Chat"):
+            chat_fast = gr.ChatInterface(
+                fn=generate_response_fast,
+                type="messages",
+                title="Quick Response",
+                description="Get faster responses without streaming",
+                examples=[
+                    "What is 15% of 240?",
+                    "Explain Newton's first law of motion",
+                    "How do you solve quadratic equations?",
+                    "What is the Pythagorean theorem?"
+                ]
+            )
+        
+        with gr.Tab("Settings"):
+            gr.Markdown("### Generation Parameters")
             max_tokens = gr.Slider(
-                minimum=50, 
-                maximum=1024, 
-                value=256, 
-                step=50, 
+                minimum=64,
+                maximum=2048,
+                value=1024,
+                step=64,
                 label="Max Tokens"
             )
             temperature = gr.Slider(
-                minimum=0.1, 
-                maximum=2.0, 
-                value=0.7, 
-                step=0.1, 
+                minimum=0.1,
+                maximum=2.0,
+                value=0.7,
+                step=0.1,
                 label="Temperature"
             )
-            
-            gr.Markdown("### Examples")
-            gr.Examples(
-                examples=[
-                    "Solve: 2x + 5 = 15",
-                    "Explain quantum mechanics simply",
-                    "What is the derivative of x²?",
-                ],
-                inputs=msg
+            top_p = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.9,
+                step=0.05,
+                label="Top P"
             )
     
-    # Chat history state
-    history = gr.State([])
-    
-    # Event handlers
-    def user_submit(message, hist):
-        return hist + [[message, None]], hist + [[message, None]], ""
-    
-    def bot_respond(hist, max_tok, temp):
-        if hist and hist[-1][1] is None:
-            message = hist[-1][0]
-            _, updated_hist, _ = generate_response(message, hist[:-1], max_tok, temp)
-            return updated_hist, updated_hist
-        return hist, hist
-    
-    # Submit message
-    msg.submit(
-        user_submit, 
-        [msg, history], 
-        [chatbot, history, msg]
-    ).then(
-        bot_respond,
-        [history, max_tokens, temperature],
-        [chatbot, history]
-    )
-    
-    send_btn.click(
-        user_submit, 
-        [msg, history], 
-        [chatbot, history, msg]
-    ).then(
-        bot_respond,
-        [history, max_tokens, temperature],
-        [chatbot, history]
-    )
-    
-    # Clear chat
-    clear_btn.click(
-        lambda: ([], []), 
-        outputs=[chatbot, history]
-    )
+    return demo
 
+# Launch the interface
 if __name__ == "__main__":
+    demo = create_interface()
     demo.launch()