Enhanced monitoring and tweaked for GPU

main.py

@@ -5,213 +5,389 @@ import json
 import os
 import psutil
 import GPUtil
+import torch
+from typing import Dict, Tuple, Any, Optional
+import logging
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('logs/app.log'),
+        logging.StreamHandler()
+    ]
+)
 
 # Create logs directory if it doesn't exist
 if not os.path.exists('logs'):
     os.makedirs('logs')
 
-class SystemMonitor:
+class EnhancedSystemMonitor:
+    """Enhanced system monitoring with detailed GPU stats and change tracking"""
     def __init__(self):
         self.gpu = None
+        self.previous_stats = None
         try:
             self.gpu = GPUtil.getGPUs()[0]
-        except:
-            pass
+            logging.info(f"GPU initialized: {self.gpu.name}")
+        except Exception as e:
+            logging.warning(f"Could not initialize GPU monitoring: {e}")
 
-    def get_stats(self):
+    def get_stats(self) -> Dict[str, float]:
+        """Get comprehensive system statistics including GPU metrics"""
         stats = {
             'cpu_percent': psutil.cpu_percent(),
             'memory_percent': psutil.virtual_memory().percent,
             'gpu_util': 0,
-            'gpu_memory': 0
+            'gpu_memory': 0,
+            'gpu_clock': 0,
+            'gpu_temp': 0,
+            'gpu_power': 0
         }
+        
         if self.gpu:
-            stats['gpu_util'] = self.gpu.load * 100
-            stats['gpu_memory'] = self.gpu.memoryUtil * 100
+            try:
+                self.gpu = GPUtil.getGPUs()[0]  # Refresh GPU info
+                stats.update({
+                    'gpu_util': self.gpu.load * 100,
+                    'gpu_memory': self.gpu.memoryUtil * 100,
+                    'gpu_clock': getattr(self.gpu, 'memoryTotal', 0),  # Using memoryTotal as a fallback
+                    'gpu_temp': getattr(self.gpu, 'temperature', 0),
+                    'gpu_power': getattr(self.gpu, 'powerDraw', 0)
+                })
+            except Exception as e:
+                logging.error(f"Error updating GPU stats: {e}")
+
+        if self.previous_stats:
+            stats['gpu_util_change'] = stats['gpu_util'] - self.previous_stats['gpu_util']
+            stats['gpu_memory_change'] = stats['gpu_memory'] - self.previous_stats['gpu_memory']
+        
+        self.previous_stats = stats.copy()
         return stats
 
 class SessionLogger:
+    """Enhanced session logging with system metrics"""
     def __init__(self):
         self.session_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         self.log_file = f"logs/session_{self.session_id}.json"
         self.session_data = {
             "session_id": self.session_id,
             "start_time": datetime.datetime.now().isoformat(),
+            "system_info": self._get_system_info(),
             "conversations": []
         }
 
-    def log_interaction(self, user_message, ai_response, stats):
+    def _get_system_info(self) -> Dict[str, Any]:
+        """Collect system information at session start"""
+        info = {
+            "cpu_freq": psutil.cpu_freq()._asdict() if psutil.cpu_freq() else {},
+            "cpu_count": psutil.cpu_count(),
+            "cpu_count_logical": psutil.cpu_count(logical=True),
+            "memory_total": psutil.virtual_memory().total,
+            "platform": os.name
+        }
+        try:
+            gpu = GPUtil.getGPUs()[0]
+            info["gpu_info"] = {
+                "name": gpu.name,
+                "memory_total": gpu.memoryTotal,
+                "driver": gpu.driver
+            }
+        except:
+            info["gpu_info"] = None
+        return info
+
+    def log_interaction(self, 
+                       user_message: str, 
+                       ai_response: str, 
+                       stats: Dict[str, float],
+                       metadata: Optional[Dict[str, Any]] = None):
+        """Log interaction with enhanced metadata"""
         interaction = {
             "timestamp": datetime.datetime.now().isoformat(),
             "user_message": user_message,
             "ai_response": ai_response,
-            "system_stats": stats
+            "system_stats": stats,
+            "metadata": metadata or {}
         }
         self.session_data["conversations"].append(interaction)
         self.save_session()
 
     def save_session(self):
-        with open(self.log_file, 'w', encoding='utf-8') as f:
-            json.dump(self.session_data, f, indent=2, ensure_ascii=False)
-
-# Initialize monitors
-system_monitor = SystemMonitor()
-session_logger = SessionLogger()
-
-# Initialize model
-MODEL_PATH = "D:/llama-tutor/models/llama-2-7b-chat.gguf"
-llm = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
-    model_type='llama',
-    context_length=2048,
-    gpu_layers=32,
-    threads=6,
-    batch_size=256
-)
+        """Save session data to file with error handling"""
+        try:
+            with open(self.log_file, 'w', encoding='utf-8') as f:
+                json.dump(self.session_data, f, indent=2, ensure_ascii=False)
+        except Exception as e:
+            logging.error(f"Error saving session data: {e}")
 
-SYSTEM_PROMPT = """You are an expert AI tutor, designed to help students learn and master any subject through engaging, Socratic dialogue. Your approach is guided by these core principles:
-
-1. TEACHING APPROACH:
-- Focus on building deep understanding rather than providing direct answers
-- Use the Socratic method: guide students through questions and critical thinking
-- Adapt your explanations to match the student's current level of understanding
-- Be patient, encouraging, and maintain a supportive learning environment"""
-
-def manage_conversation_length(history, max_length=4000):
-    """Manage conversation history to prevent context overflow"""
-    current_length = 0
-    truncated_history = []
-
-    # Start from the most recent messages
-    for h in reversed(history):
-        # Estimate tokens (roughly 4 chars per token)
-        message_length = len(h[0]) + len(h[1])
-        if current_length + message_length < max_length:
-            truncated_history.insert(0, h)
-            current_length += message_length
-        else:
-            # If including the next message would exceed the max length,
-            # truncate the current message to fit within the remaining space
-            remaining_space = max_length - current_length
-            if remaining_space > 0:
-                truncated_message = h[1][:remaining_space]
-                truncated_history.insert(0, [h[0], truncated_message])
-            break
-
-    return truncated_history
-
-def generate_response(message, history):
-    # Manage conversation length
-    managed_history = manage_conversation_length(history)
-    
-    # Format the conversation history
-    chat_history = "\n".join([f"Human: {h}\nAssistant: {a}\n" for h, a in managed_history])
-    
-    # Calculate approximate token length
-    total_length = len(SYSTEM_PROMPT) + len(chat_history) + len(message)
-    approximate_tokens = total_length // 4  # Rough estimation
-    
-    prompt = f"""<s>[INST] <<SYS>>{SYSTEM_PROMPT}<</SYS>>
+class ModelManager:
+    """Manages model initialization, inference, and resource handling"""
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.llm = None
+        self.config = None
+        self.initialize_model()
+        self.total_tokens_processed = 0
 
-Previous conversation:
-{chat_history}
+    def initialize_model(self, gpu_layers: int = 24, batch_size: int = 128) -> None:
+        """Initialize the model with specified parameters"""
+        try:
+            self.llm = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                model_type='llama',
+                context_length=2048,
+                gpu_layers=gpu_layers,
+                threads=4,
+                batch_size=batch_size
+            )
+            self.config = {'gpu_layers': gpu_layers, 'batch_size': batch_size}
+            logging.info(f"Model initialized with config: {self.config}")
+        except Exception as e:
+            logging.error(f"Error initializing model: {e}")
+            raise
 
-Current question: {message} [/INST]"""
-    
-    # Get system stats before processing
-    stats_before = system_monitor.get_stats()
-    
-    # Update stats with token information
-    stats_before['approximate_tokens'] = approximate_tokens
-    
-    response = llm(prompt, 
-              max_new_tokens=1024,  # Increased from 512
-              temperature=0.7,
-              top_p=0.9,
-              top_k=40,
-              stop=["</s>"])
-    
-    # Log interaction with system stats
-    session_logger.log_interaction(message, response, stats_before)
-    
-    return response
+    def generate_with_monitoring(self, 
+                               prompt: str, 
+                               system_monitor: EnhancedSystemMonitor,
+                               **kwargs) -> Tuple[str, Dict[str, Any]]:
+        """Generate response with comprehensive monitoring"""
+        pre_stats = system_monitor.get_stats()
+        
+        try:
+            # Generate response
+            response = self.llm(prompt, **kwargs)
+            post_stats = system_monitor.get_stats()
+            
+            # Calculate tokens for this interaction
+            input_tokens = len(prompt) // 4
+            output_tokens = len(response) // 4
+            total_tokens = input_tokens + output_tokens
+            
+            # Update total tokens processed
+            self.total_tokens_processed += total_tokens
+            
+            # Update stats with token information
+            post_stats.update({
+                'approximate_tokens': total_tokens,
+                'total_tokens_processed': self.total_tokens_processed,
+                'input_tokens': input_tokens,
+                'output_tokens': output_tokens
+            })
+            
+            return response, {
+                'pre_stats': pre_stats,
+                'post_stats': post_stats,
+                'tokens_processed': total_tokens,
+                'total_tokens': self.total_tokens_processed
+            }
+            
+        except RuntimeError as e:
+            if "out of memory" in str(e):
+                logging.warning("OOM detected, attempting recovery...")
+                torch.cuda.empty_cache()
+                
+                # Reduce parameters and retry
+                new_config = {
+                    'gpu_layers': max(8, self.config['gpu_layers'] - 4),
+                    'batch_size': max(32, self.config['batch_size'] // 2)
+                }
+                
+                self.initialize_model(**new_config)
+                return self.generate_with_monitoring(prompt, system_monitor, **kwargs)
+            raise
 
-def update_stats():
-    stats = system_monitor.get_stats()
-    current_tokens = len(SYSTEM_PROMPT) // 4  # Base system prompt tokens
+def format_stats(stats: Dict[str, float]) -> str:
+    """Format system statistics for display"""
     return f"""System Statistics:
-    CPU Usage: {stats['cpu_percent']:.1f}%
-    Memory Usage: {stats['memory_percent']:.1f}%
-    GPU Utilization: {stats['gpu_util']:.1f}%
-    GPU Memory: {stats['gpu_memory']:.1f}%
-    Context Usage: ~{current_tokens} tokens"""
-
-# Create the Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎓 TA Tutor")
+    CPU Usage: {stats.get('cpu_percent', 0):.1f}%
+    Memory Usage: {stats.get('memory_percent', 0):.1f}%
     
-    with gr.Row():
-        with gr.Column(scale=4):
-            chatbot = gr.Chatbot(
-                value=[],
-                label="Tutoring Session",
-                height=380
-            )
-            with gr.Row():
-                msg = gr.Textbox(
-                    label="Your Question",
-                    placeholder="Ask your question here...",
-                    lines=2
+    GPU Status:
+    • Utilization: {stats.get('gpu_util', 0):.1f}% ({stats.get('gpu_util_change', 0):.1f}% change)
+    • Memory: {stats.get('gpu_memory', 0):.1f}% ({stats.get('gpu_memory_change', 0):.1f}% change)
+    • Memory Total: {stats.get('gpu_clock', 0)} MB
+    • Temperature: {stats.get('gpu_temp', 0)}°C
+    
+    Model Info:
+    • Current Usage: ~{stats.get('approximate_tokens', 0)} tokens
+    • Input Tokens: ~{stats.get('input_tokens', 0)}
+    • Output Tokens: ~{stats.get('output_tokens', 0)}
+    • Total Processed: {stats.get('total_tokens_processed', 0)} tokens"""
+
+def create_demo(model_path: str):
+    """Create and configure the Gradio demo"""
+    # Initialize components
+    system_monitor = EnhancedSystemMonitor()
+    session_logger = SessionLogger()
+    model_manager = ModelManager(model_path)
+
+    # System prompt
+    SYSTEM_PROMPT = """You are an expert AI tutor, designed to help students learn and master any subject through engaging, Socratic dialogue. Your approach is guided by these core principles:
+
+    1. TEACHING APPROACH:
+    - Focus on building deep understanding rather than providing direct answers
+    - Use the Socratic method: guide students through questions and critical thinking
+    - Adapt your explanations to match the student's current level of understanding
+    - Be patient, encouraging, and maintain a supportive learning environment"""
+
+    def manage_conversation_length(history, max_length: int = 4000) -> list:
+        """Manage conversation history to prevent context overflow"""
+        current_length = 0
+        truncated_history = []
+        
+        for h in reversed(history):
+            message_length = len(h['content'])
+            if current_length + message_length < max_length:
+                truncated_history.insert(0, h)
+                current_length += message_length
+            else:
+                remaining_space = max_length - current_length
+                if remaining_space > 0:
+                    truncated_message = h.copy()
+                    truncated_message['content'] = h['content'][:remaining_space]
+                    truncated_history.insert(0, truncated_message)
+                break
+        
+        return truncated_history
+
+    def generate_response(message: str, history: list) -> Tuple[str, Dict[str, Any]]:
+        """Generate response with conversation management"""
+        managed_history = manage_conversation_length(history)
+        
+        chat_history = "\n".join([
+            f"{'Human' if h['role'] == 'user' else 'Assistant'}: {h['content']}"
+            for h in managed_history
+        ])
+        
+        prompt = f"""<s>[INST] <<SYS>>{SYSTEM_PROMPT}<</SYS>>
+
+        Previous conversation:
+        {chat_history}
+
+        Current question: {message} [/INST]"""
+        
+        # Calculate approximate tokens for the full context
+        system_tokens = len(SYSTEM_PROMPT) // 4
+        history_tokens = len(chat_history) // 4
+        message_tokens = len(message) // 4
+        total_input_tokens = system_tokens + history_tokens + message_tokens
+        
+        # Update system monitor with token count
+        current_stats = system_monitor.get_stats()
+        current_stats['approximate_tokens'] = total_input_tokens
+        
+        response, monitoring_data = model_manager.generate_with_monitoring(
+            prompt,
+            system_monitor,
+            max_new_tokens=1024,
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            stop=["</s>"]
+        )
+        
+        # Update monitoring data with token information
+        monitoring_data['post_stats']['approximate_tokens'] = total_input_tokens + (len(response) // 4)
+        
+        session_logger.log_interaction(
+            message,
+            response,
+            monitoring_data['post_stats'],
+            {"monitoring_data": monitoring_data}
+        )
+        
+        return response, monitoring_data
+
+    # Create the Gradio interface
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🎓 TA Tutor")
+        
+        with gr.Row():
+            with gr.Column(scale=4):
+                chatbot = gr.Chatbot(
+                    value=[],
+                    label="Tutoring Session",
+                    height=380,
+                    type="messages"  # Updated format
+                )
+                
+                with gr.Row():
+                    msg = gr.Textbox(
+                        label="Your Question",
+                        placeholder="Ask your question here...",
+                        lines=2
+                    )
+                    submit = gr.Button("Send", variant="primary")
+                
+                # Example questions
+                example_questions = [
+                    "Can you help me understand quantum mechanics?",
+                    "What's the best way to learn calculus?",
+                    "Explain photosynthesis in simple terms.",
+                    "Can you create a study plan for me?",
+                    "A car accelerates uniformly from rest to a speed of 20 m/s in 5 seconds. What is the acceleration of the car?",
+                    "Balance the following chemical equation: Fe + O2 -> Fe2O3",
+                    "Find the derivative of the function f(x) = 3x^4 - 2x^3 + 5x - 7.",
+                    "Given the matrix A = [[1, 2], [3, 4]], find the determinant of A.",
+                    "Write a Python function that takes a list of integers as input and returns the sum of all even numbers in the list.",
+                    "Two forces, F1 = 20 N and F2 = 30 N, act on a point at an angle of 60° to each other. Find the magnitude and direction of the resultant force.",
+                    "In a series circuit with a 12V battery, there are two resistors: R1 = 4Ω and R2 = 6Ω. Find the current flowing through the circuit.",
+                    "Describe the basic steps in the engineering design process.",
+                    "Explain the difference between elastic and plastic deformation in materials.",
+                    "State the first law of thermodynamics and provide an example of its application in engineering.",
+                ]
+                gr.Examples(
+                    examples=example_questions,
+                    inputs=msg,
+                    label="Example Questions"
                 )
-                submit = gr.Button("Send", variant="primary")
             
-            example_questions = [
-                "Can you help me understand quantum mechanics?",
-                "What's the best way to learn calculus?",
-                "Explain photosynthesis in simple terms.",
-                "Can you create a study plan for me?",
-                "A car accelerates uniformly from rest to a speed of 20 m/s in 5 seconds. What is the acceleration of the car?",
-                "Balance the following chemical equation: Fe + O2 -> Fe2O3",
-                "Find the derivative of the function f(x) = 3x^4 - 2x^3 + 5x - 7.",
-                "Given the matrix A = [[1, 2], [3, 4]], find the determinant of A.",
-                "Write a Python function that takes a list of integers as input and returns the sum of all even numbers in the list.",
-                "Two forces, F1 = 20 N and F2 = 30 N, act on a point at an angle of 60° to each other. Find the magnitude and direction of the resultant force.",
-                "In a series circuit with a 12V battery, there are two resistors: R1 = 4Ω and R2 = 6Ω. Find the current flowing through the circuit.",
-                "Describe the basic steps in the engineering design process.",
-                "Explain the difference between elastic and plastic deformation in materials.",
-                "State the first law of thermodynamics and provide an example of its application in engineering."
-            ]
-            gr.Examples(
-                examples=example_questions,
-                inputs=msg,
-                label="Example Questions"
-            )
+            with gr.Column(scale=1):
+                stats_display = gr.Textbox(
+                    label="System Performance",
+                    lines=8,
+                    interactive=False,
+                    value=format_stats(system_monitor.get_stats())
+                )
         
-        with gr.Column(scale=1):
-            stats = gr.Textbox(
-                label="System Performance",
-                lines=6,
-                interactive=False,
-                value=update_stats()
-            )
+        def respond(message: str, history: list) -> Tuple[str, list, str]:
+            """Handle user input and generate response"""
+            try:
+                response, monitoring_data = generate_response(message, history)
+                
+                history.append({"role": "user", "content": message})
+                history.append({"role": "assistant", "content": response})
+                
+                new_stats = format_stats(monitoring_data['post_stats'])
+                return "", history, new_stats
+                
+            except Exception as e:
+                error_message = f"Error during generation: {str(e)}"
+                logging.error(error_message)
+                history.append({"role": "system", "content": error_message})
+                return "", history, format_stats(system_monitor.get_stats())
 
-    def respond(message, history):
-        bot_message = generate_response(message, history)
-        history.append([message, bot_message])
-        new_stats = update_stats()
-        return "", history, new_stats
+        # Connect UI elements
+        msg.submit(respond, [msg, chatbot], [msg, chatbot, stats_display])
+        submit.click(respond, [msg, chatbot], [msg, chatbot, stats_display])
 
-    # Connect message input and submit button
-    msg.submit(respond, [msg, chatbot], [msg, chatbot, stats])
-    submit.click(respond, [msg, chatbot], [msg, chatbot, stats])
+        # Add stats refresh button
+        refresh = gr.Button("🔄 Refresh Stats")
+        refresh.click(
+            lambda: format_stats(system_monitor.get_stats()),
+            outputs=[stats_display]
+        )
 
-    # Add a refresh button for stats
-    refresh = gr.Button("🔄 Refresh Stats")
-    refresh.click(update_stats, outputs=[stats])
+    return demo
 
 if __name__ == "__main__":
+    MODEL_PATH = "D:/llama-tutor/models/llama-2-7b-chat.gguf"
+    demo = create_demo(MODEL_PATH)
     demo.queue()
     demo.launch(
         server_port=7860,
-        server_name="0.0.0.0"
+        server_name="127.0.0.1",  # Changed to localhost for security
+        show_error=True
     )