Update app.py

app.py

@@ -1,29 +1,34 @@
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
+from huggingface_hub import login
 import gradio as gr
 import torch
 import os
+import time
 
-# ===== AUTO-DEVICE CONFIGURATION =====
+# ===== AUTHENTICATION =====
+login(token=os.environ.get("HF_TOKEN", "your_hf_token_here"))  # Set in Space secrets
+
+# ===== SMART DEVICE CONFIGURATION =====
 def configure_device():
-    """Smart device selection with performance optimizations"""
+    """Auto-configure GPU/CPU with optimizations"""
     if torch.cuda.is_available():
-        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-        torch.backends.cudnn.benchmark = True  # Auto-tunes CUDA
-        return 0, torch.float16  # GPU with half-precision
-    
-    # Optimized CPU configuration
-    torch.set_num_threads(min(4, os.cpu_count() or 1))
-    return -1, torch.float32
+        torch.backends.cudnn.benchmark = True
+        return 0, torch.float16  # GPU with half precision
+    torch.set_num_threads(os.cpu_count() or 4)
+    return -1, torch.float32  # CPU fallback
 
 device, dtype = configure_device()
-device_name = "GPU: " + torch.cuda.get_device_name(0) if device == 0 else "CPU"
+device_name = f"{'GPU: ' + torch.cuda.get_device_name(0) if device == 0 else 'CPU'}"
 print(f"⚡ Running on: {device_name} | Precision: {dtype}")
 
-# ===== BULLETPROOF MODEL LOADING =====
+# ===== MODEL LOADING =====
+MODEL_NAME = "google/gemma-2b-it"
 try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = pipeline(
         task="text-generation",
-        model="google/gemma-2b-it",  # Fast 2B parameter model
+        model=MODEL_NAME,
+        tokenizer=tokenizer,
         device=device,
         torch_dtype=dtype,
         model_kwargs={
@@ -31,51 +36,105 @@ try:
             "trust_remote_code": True
         }
     )
-    
-    # Pre-warm model (critical for fast first response)
-    model("Warming up...", max_new_tokens=1)
-    
+    # Warm-up
+    model("Explain 2+2:", max_new_tokens=10)
 except Exception as e:
-    # Fallback to CPU if GPU fails
-    print(f"⚠️ GPU failed, falling back to CPU: {str(e)}")
-    device, dtype = -1, torch.float32
-    model = pipeline(
-        task="text-generation",
-        model="google/gemma-2b-it",
-        device=device,
-        torch_dtype=dtype
-    )
+    raise gr.Error(f"""🚨 Model loading failed. Ensure:
+1. You accepted Gemma's terms at: https://huggingface.co/{MODEL_NAME}
+2. HF_TOKEN is set in Space secrets
+Error: {str(e)}""")
 
-# ===== ULTRA-FAST GENERATION =====
-def generate(prompt):
-    """Guaranteed fast response (1-2 seconds)"""
+# ===== ENHANCED GENERATION =====
+def generate_with_explanation(prompt):
+    """Generate step-by-step explanations with performance monitoring"""
+    start_time = time.time()
+    
     try:
-        return model(
-            prompt,
-            max_new_tokens=50,  # Optimal for speed
-            temperature=0.1,    # More deterministic
-            do_sample=False,    # Disable sampling for speed
-            pad_token_id=model.tokenizer.eos_token_id
-        )[0]['generated_text']
+        # Enhanced prompt for step-by-step answers
+        enhanced_prompt = f"""Explain step-by-step in detail:
+
+Question: {prompt}
+
+Answer:"""
+        
+        # Generation with optimized parameters
+        output = model(
+            enhanced_prompt,
+            max_new_tokens=150,
+            temperature=0.3,
+            top_k=50,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        # Extract and format response
+        full_response = output[0]['generated_text']
+        answer = full_response.split("Answer:")[-1].strip()
+        
+        # Performance metrics
+        gen_time = time.time() - start_time
+        tokens_sec = len(answer.split()) / gen_time
+        
+        return f"""{answer}
+
+⏱️ Generated in {gen_time:.2f}s ({tokens_sec:.1f} tokens/sec)"""
+    
     except Exception as e:
-        return f"🔴 Error (but UI keeps working): {str(e)}"
+        return f"❌ Error: {str(e)}"
 
-# ===== LIGHTNING-FAST INTERFACE =====
-with gr.Blocks(title="⚡ Instant AI (1-2s responses)") as demo:
-    gr.Markdown("## Type anything for instant answers:")
+# ===== ADVANCED INTERFACE =====
+with gr.Blocks(theme=gr.themes.Soft(), title="🧠 AI Tutor with Step-by-Step Explanations") as demo:
+    # Header
+    gr.Markdown("""<h1><center>Step-by-Step AI Tutor</center></h1>""")
+    
+    # Input Section
     with gr.Row():
-        inp = gr.Textbox(placeholder="How does photosynthesis work?", 
-                        lines=2,
-                        max_lines=3)
+        input_box = gr.Textbox(
+            label="Ask anything",
+            placeholder="E.g. 'Explain quantum computing basics'",
+            lines=3,
+            max_lines=5
+        )
+    
+    # Control Panel
+    with gr.Row():
+        submit_btn = gr.Button("Generate Explanation", variant="primary")
+        clear_btn = gr.Button("Clear")
+    
+    # Output Section
     with gr.Row():
-        out = gr.Textbox(label="Answer appears here (1-2 seconds)",
-                        lines=5)
-    inp.submit(generate, inp, out)
+        output_box = gr.Textbox(
+            label="Detailed Explanation",
+            lines=10,
+            interactive=False
+        )
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            "Explain how photosynthesis works step by step",
+            "Solve 3x + 5 = 20 showing each step",
+            "Describe the water cycle with bullet points"
+        ],
+        inputs=input_box
+    )
+    
+    # Event Handlers
+    submit_btn.click(
+        fn=generate_with_explanation,
+        inputs=input_box,
+        outputs=output_box
+    )
+    clear_btn.click(
+        fn=lambda: ("", ""),
+        inputs=None,
+        outputs=[input_box, output_box]
+    )
 
-# ===== FAILSAFE LAUNCH =====
+# ===== LAUNCH =====
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        show_error=True
+        share=False
     )