Update app.py

app.py

@@ -2,86 +2,153 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
+import time
+import logging
+from datetime import datetime
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    datefmt='%H:%M:%S'
+)
+logger = logging.getLogger(__name__)
 
 # Configuration for CPU optimization
 class Config:
     MODEL_PATH = "navidfalah/3ai"
     BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.1"
-    MAX_NEW_TOKENS = 150  # Much shorter for faster generation
+    MAX_NEW_TOKENS = 50  # Very short for speed
     TEMPERATURE = 0.7
     TOP_P = 0.9
-    MAX_INPUT_LENGTH = 256  # Shorter input for faster processing
+    MAX_INPUT_LENGTH = 128  # Very short input
+    USE_SMALL_MODEL = True  # Set to True to use tiny model
 
 # Global variables
 model = None
 tokenizer = None
+model_load_time = None
+
+def log_time(start_time, operation):
+    """Log time taken for an operation."""
+    elapsed = time.time() - start_time
+    logger.info(f"{operation} took {elapsed:.2f} seconds")
+    return elapsed
 
 def load_model_cpu_optimized():
-    """Load model optimized for CPU inference."""
-    global model, tokenizer
+    """Load model optimized for CPU inference with timing."""
+    global model, tokenizer, model_load_time
     
     if model is not None and tokenizer is not None:
+        logger.info("Model already loaded, using cached version")
         return model, tokenizer
     
+    total_start = time.time()
+    
     try:
-        print("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(Config.BASE_MODEL)
-        if tokenizer.pad_token is None:
+        # Use smaller model for CPU
+        if Config.USE_SMALL_MODEL:
+            logger.info("Loading small model for fast CPU inference...")
+            
+            # Option 1: Use GPT-2 small (fastest)
+            model_name = "gpt2"  # Only 124M parameters
+            
+            # Option 2: Use DistilGPT-2 (even faster)
+            # model_name = "distilgpt2"  # Only 82M parameters
+            
+            start = time.time()
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
             tokenizer.pad_token = tokenizer.eos_token
+            log_time(start, "Tokenizer loading")
+            
+            start = time.time()
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float32,
+                low_cpu_mem_usage=True
+            )
+            model.eval()
+            log_time(start, "Model loading")
+            
+            logger.info(f"✅ Loaded {model_name} for fast inference")
         
-        print("Loading model for CPU...")
-        # Load in float32 for CPU (no quantization)
-        model = AutoModelForCausalLM.from_pretrained(
-            Config.BASE_MODEL,
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            low_cpu_mem_usage=True,
-            device_map="cpu"  # Force CPU
-        )
+        else:
+            # Load full model (will be slow on CPU)
+            logger.warning("Loading full Mistral model - this will be VERY slow on CPU!")
+            logger.warning("Consider setting USE_SMALL_MODEL = True")
+            
+            start = time.time()
+            tokenizer = AutoTokenizer.from_pretrained(Config.BASE_MODEL)
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            log_time(start, "Tokenizer loading")
+            
+            start = time.time()
+            model = AutoModelForCausalLM.from_pretrained(
+                Config.BASE_MODEL,
+                torch_dtype=torch.float32,
+                low_cpu_mem_usage=True,
+                device_map="cpu"
+            )
+            model.eval()
+            log_time(start, "Model loading")
+        
+        model_load_time = log_time(total_start, "Total model loading")
+        logger.info(f"Model size: ~{sum(p.numel() for p in model.parameters()) / 1e6:.0f}M parameters")
         
-        model.eval()
-        print("✅ Model loaded on CPU")
         return model, tokenizer
         
     except Exception as e:
-        print(f"Error loading model: {e}")
-        # Try a smaller model as fallback
-        try:
-            print("Trying smaller model fallback...")
-            model = AutoModelForCausalLM.from_pretrained(
-                "gpt2",  # Much smaller fallback model
-                torch_dtype=torch.float32
-            )
-            tokenizer = AutoTokenizer.from_pretrained("gpt2")
-            tokenizer.pad_token = tokenizer.eos_token
-            model.eval()
-            print("✅ Loaded fallback model (GPT-2)")
-            return model, tokenizer
-        except:
-            return None, None
+        logger.error(f"Failed to load model: {e}")
+        return None, None
 
-def analyze_text(user_input):
-    """Simple and fast text analysis."""
+def analyze_text(user_input, progress=gr.Progress()):
+    """Simple and fast text analysis with progress tracking."""
+    start_time = time.time()
+    
     if not user_input.strip():
-        return "Please enter some text to analyze."
+        return "Please enter some text to analyze.", "No input provided"
+    
+    logger.info(f"Starting analysis for input: {user_input[:50]}...")
+    
+    # Update progress
+    progress(0.1, desc="Loading model...")
     
+    # Load model with timing
+    model_start = time.time()
     model, tokenizer = load_model_cpu_optimized()
+    model_time = time.time() - model_start
     
     if model is None or tokenizer is None:
-        return "Error: Could not load model. Please try again."
+        return "Error: Could not load model.", f"Model loading failed after {model_time:.2f}s"
+    
+    progress(0.3, desc="Model loaded, preparing input...")
     
     try:
-        # Simple prompt - no complex formatting
-        prompt = f"Analyze this life situation and provide brief advice: {user_input}\n\nAnalysis:"
+        # Simple prompt
+        prompt = f"Life advice for: {user_input}\n\nAdvice:"
+        logger.info(f"Prompt length: {len(prompt)} characters")
         
-        # Tokenize with minimal length
+        # Tokenize with timing
+        tokenize_start = time.time()
         inputs = tokenizer(
             prompt, 
             return_tensors="pt", 
             truncation=True, 
             max_length=Config.MAX_INPUT_LENGTH
         )
+        tokenize_time = log_time(tokenize_start, "Tokenization")
+        
+        progress(0.5, desc="Generating response...")
+        
+        # Log input details
+        input_ids = inputs['input_ids']
+        logger.info(f"Input tokens: {input_ids.shape[1]}")
+        
+        # Generate with timing
+        gen_start = time.time()
+        logger.info(f"Starting generation with max {Config.MAX_NEW_TOKENS} new tokens...")
         
-        # Generate with aggressive settings for speed
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
@@ -89,61 +156,129 @@ def analyze_text(user_input):
                 temperature=Config.TEMPERATURE,
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
-                early_stopping=True,  # Stop as soon as possible
-                num_beams=1  # No beam search for speed
+                early_stopping=True,
+                num_beams=1,
+                use_cache=True  # Enable KV cache
             )
         
+        gen_time = log_time(gen_start, "Generation")
+        logger.info(f"Generated {outputs.shape[1] - input_ids.shape[1]} new tokens")
+        
+        progress(0.8, desc="Decoding response...")
+        
+        # Decode with timing
+        decode_start = time.time()
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        decode_time = log_time(decode_start, "Decoding")
         
-        # Extract only the generated part
+        # Extract generated part
         result = response[len(prompt):].strip()
         
         if not result:
-            result = "Analysis: Based on your input, I recommend focusing on balance and gradual improvements."
+            result = "Based on your input, I recommend focusing on balance and gradual improvements."
+        
+        # Total time
+        total_time = time.time() - start_time
+        logger.info(f"Total analysis time: {total_time:.2f}s")
+        
+        # Create timing report
+        timing_report = f"""### Timing Report
+- Model Load: {model_time:.2f}s {'' if model_time < 1 else '(cached)'}
+- Tokenization: {tokenize_time:.2f}s
+- Generation: {gen_time:.2f}s ({Config.MAX_NEW_TOKENS} tokens)
+- Decoding: {decode_time:.2f}s
+- **Total: {total_time:.2f}s**
+
+Model: {model.__class__.__name__}
+Input tokens: {input_ids.shape[1]}
+Output tokens: {outputs.shape[1] - input_ids.shape[1]}
+"""
         
-        return result
+        progress(1.0, desc="Complete!")
+        
+        return result, timing_report
         
     except Exception as e:
-        return f"Error: {str(e)}"
+        error_msg = f"Error during analysis: {str(e)}"
+        logger.error(error_msg)
+        total_time = time.time() - start_time
+        return error_msg, f"Failed after {total_time:.2f}s\nError: {str(e)}"
 
-# Simple Gradio Interface
-with gr.Blocks(title="Quick Life Analysis", css="footer {display: none !important}") as demo:
-    gr.Markdown("# Quick Life Satisfaction Analysis")
-    gr.Markdown("Enter your situation and get instant AI advice (optimized for CPU)")
+# Simple Gradio Interface with timing display
+with gr.Blocks(title="Fast CPU Analysis", theme=gr.themes.Base()) as demo:
+    gr.Markdown("""
+    # Fast Life Analysis (CPU Optimized)
+    
+    ⚡ Using small model for fast CPU inference. First run will be slower due to model loading.
+    """)
     
     with gr.Row():
         with gr.Column():
             input_text = gr.Textbox(
                 label="Your Input",
-                placeholder="Example: I'm stressed at work (3/10) but happy with family (8/10)...",
-                lines=4
+                placeholder="Describe your situation briefly...",
+                lines=3
             )
-            submit_btn = gr.Button("Analyze", variant="primary")
+            
+            with gr.Row():
+                submit_btn = gr.Button("🚀 Analyze (Fast)", variant="primary")
+                clear_btn = gr.Button("Clear")
+            
+            # Model selection
+            gr.Markdown("""
+            **Tips for faster response:**
+            - Keep input short (< 50 words)
+            - First analysis is slowest (model loading)
+            - Subsequent analyses are much faster
+            """)
         
         with gr.Column():
             output_text = gr.Textbox(
                 label="AI Analysis",
-                lines=6,
+                lines=4,
                 interactive=False
             )
+            timing_info = gr.Markdown(
+                label="Performance Metrics",
+                value="*Timing information will appear here*"
+            )
     
-    # Simple examples
+    # Examples
     gr.Examples(
         examples=[
-            "Work stress is high, health is okay, finances tight",
-            "Happy with job but no work-life balance",
-            "Good health and relationships but career is stagnant"
+            "Stressed at work, need balance",
+            "Happy but financially worried",
+            "Good job, poor health"
         ],
-        inputs=input_text
+        inputs=input_text,
+        label="Quick Examples (short = faster)"
     )
     
+    # Event handlers
     submit_btn.click(
         fn=analyze_text,
         inputs=input_text,
-        outputs=output_text
+        outputs=[output_text, timing_info]
+    )
+    
+    clear_btn.click(
+        fn=lambda: ("", "", "*Timing information will appear here*"),
+        outputs=[input_text, output_text, timing_info]
     )
 
 if __name__ == "__main__":
-    print("Starting CPU-optimized app...")
-    print("Note: First generation will be slow due to model loading")
+    logger.info("="*50)
+    logger.info("Starting Fast CPU-Optimized App")
+    logger.info(f"PyTorch version: {torch.__version__}")
+    logger.info(f"Using small model: {Config.USE_SMALL_MODEL}")
+    logger.info(f"Max new tokens: {Config.MAX_NEW_TOKENS}")
+    logger.info("="*50)
+    
+    # Pre-load model for faster first inference
+    logger.info("Pre-loading model...")
+    pre_load_start = time.time()
+    load_model_cpu_optimized()
+    logger.info(f"Model pre-loaded in {time.time() - pre_load_start:.2f}s")
+    
+    demo.queue()
     demo.launch()