Update app.py

app.py

@@ -3,6 +3,9 @@ import torch
 import sys
 import os
 import re
+import json
+import time
+from datetime import datetime
 from pathlib import Path
 
 # Add the project root to Python path
@@ -11,12 +14,58 @@ sys.path.append(str(project_root))
 
 from src.inference.inference import tokenizer, model  # Import from your inference.py
 from src.vector_db.manager import ChromaVectorDBManager
+from src.utils.performance import PerformanceMonitor
 import logging
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Performance history file
+PERFORMANCE_HISTORY_FILE = Path("performance_history.json")
+
+def save_performance_metrics(metrics_data):
+    """Save performance metrics to history file"""
+    try:
+        if PERFORMANCE_HISTORY_FILE.exists():
+            with open(PERFORMANCE_HISTORY_FILE, 'r') as f:
+                history = json.load(f)
+        else:
+            history = []
+        
+        history.append(metrics_data)
+        
+        with open(PERFORMANCE_HISTORY_FILE, 'w') as f:
+            json.dump(history, f, indent=2)
+            
+    except Exception as e:
+        logger.error(f"Failed to save performance metrics: {e}")
+
+def calculate_performance_metrics(start_time, end_time, prompt_tokens, generated_tokens, peak_memory_mb):
+    """Calculate performance metrics similar to the requested format"""
+    inference_time = end_time - start_time
+    total_tokens = prompt_tokens + generated_tokens
+    
+    # Calculate throughput (tokens per second)
+    throughput = total_tokens / inference_time if inference_time > 0 else 0
+    
+    # Calculate inference latency (time per token in milliseconds)
+    latency_ms = (inference_time * 1000) / total_tokens if total_tokens > 0 else 0
+    
+    return {
+        "timestamp": datetime.now().isoformat(),
+        "model": "Gemma-3-270M",
+        "load_time_s": "N/A",  # Model is already loaded
+        "inference_latency_ms": round(latency_ms, 2),
+        "throughput_tokens_s": round(throughput, 2),
+        "ram_usage_mb": round(peak_memory_mb, 2),
+        "vram_usage_mb": 0,  # CPU-only model
+        "energy_j": "N/A",  # Would require specialized monitoring
+        "prompt_tokens": prompt_tokens,
+        "generated_tokens": generated_tokens,
+        "total_inference_time_s": round(inference_time, 3)
+    }
+
 # Initialize Vector DB Manager
 try:
     logger.info("Initializing ChromaDB manager")
@@ -92,21 +141,30 @@ QUESTION:
 
 ANSWER:"""
         
-        # Use inference setup
+        # Count prompt tokens
+        prompt_tokens = len(tokenizer.encode(prompt))
+        
+        # Use inference setup with performance monitoring
         inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
         
-        # Generate response with conservative parameters for gemma-3-270m
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=256,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id
-            )
+        # Start performance monitoring for inference
+        with PerformanceMonitor("Model_Inference") as monitor:
+            start_time = time.time()
+            
+            # Generate response with conservative parameters for gemma-3-270m
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=256,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    repetition_penalty=1.1,
+                    pad_token_id=tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.eos_token_id
+                )
+            
+            end_time = time.time()
         
         # Decode and clean response
         full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -114,16 +172,50 @@ ANSWER:"""
         # Extract only the generated part (remove the original prompt)
         answer = full_response[len(prompt):].strip()
         
+        # Count generated tokens
+        generated_tokens = len(tokenizer.encode(answer))
+        
+        # Get performance metrics from monitor
+        perf_metrics = monitor.stop_monitoring()
+        
+        # Calculate and save performance metrics
+        metrics_data = calculate_performance_metrics(
+            start_time, 
+            end_time, 
+            prompt_tokens, 
+            generated_tokens, 
+            perf_metrics.peak_memory
+        )
+        
+        # Save to history
+        save_performance_metrics(metrics_data)
+        
+        # Log performance summary
+        logger.info(f"Performance Metrics:")
+        logger.info(f"  Model: {metrics_data['model']}")
+        logger.info(f"  Inference Latency: {metrics_data['inference_latency_ms']} ms")
+        logger.info(f"  Throughput: {metrics_data['throughput_tokens_s']} tokens/s")
+        logger.info(f"  RAM Usage: {metrics_data['ram_usage_mb']} MB")
+        logger.info(f"  Tokens (prompt/generated): {metrics_data['prompt_tokens']}/{metrics_data['generated_tokens']}")
+        
         if not answer:
             answer = "I apologize, but I couldn't generate a proper response. Please try rephrasing your question."
         
         logger.info(f"Generated response length: {len(answer)} characters")
         
+        # Add performance info to sources
+        perf_info = f"\n\n**Performance Metrics:**\n" \
+                   f"- Model: {metrics_data['model']}\n" \
+                   f"- Inference Latency: {metrics_data['inference_latency_ms']} ms\n" \
+                   f"- Throughput: {metrics_data['throughput_tokens_s']} tokens/s\n" \
+                   f"- RAM Usage: {metrics_data['ram_usage_mb']} MB\n" \
+                   f"- Total Inference Time: {metrics_data['total_inference_time_s']} s"
+        
         # Return answer and sources if requested
         if show_context:
-            return answer, f"**Sources Used:**\n{sources}\n\n**Context:**\n{context}"
+            return answer, f"**Sources Used:**\n{sources}\n\n**Context:**\n{context}{perf_info}"
         else:
-            return answer, f"**Sources Used:**\n{sources}"
+            return answer, f"**Sources Used:**\n{sources}{perf_info}"
             
     except Exception as e:
         logger.error(f"Error in chat_with_rag: {e}")