Add history and metrics persistence with file storage

app.py

@@ -25,6 +25,9 @@ MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
 CACHE_DIR = "/home/user/.cache/huggingface"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
+# History file
+HISTORY_FILE = "review_history.json"
+
 class Review:
     def __init__(self, code: str, language: str, suggestions: str):
         self.code = code
@@ -32,6 +35,22 @@ class Review:
         self.suggestions = suggestions
         self.timestamp = datetime.now().isoformat()
         self.response_time = 0.0
+    
+    def to_dict(self):
+        return {
+            'timestamp': self.timestamp,
+            'language': self.language,
+            'code': self.code,
+            'suggestions': self.suggestions,
+            'response_time': self.response_time
+        }
+    
+    @classmethod
+    def from_dict(cls, data):
+        review = cls(data['code'], data['language'], data['suggestions'])
+        review.timestamp = data['timestamp']
+        review.response_time = data['response_time']
+        return review
 
 class CodeReviewer:
     def __init__(self):
@@ -45,6 +64,32 @@ class CodeReviewer:
             'reviews_today': 0
         }
         self._initialized = False
+        self.load_history()
+        
+    def load_history(self):
+        """Load review history from file."""
+        try:
+            if os.path.exists(HISTORY_FILE):
+                with open(HISTORY_FILE, 'r') as f:
+                    data = json.load(f)
+                    self.review_history = [Review.from_dict(r) for r in data['history']]
+                    self.metrics = data['metrics']
+                logger.info(f"Loaded {len(self.review_history)} reviews from history")
+        except Exception as e:
+            logger.error(f"Error loading history: {e}")
+    
+    def save_history(self):
+        """Save review history to file."""
+        try:
+            data = {
+                'history': [r.to_dict() for r in self.review_history],
+                'metrics': self.metrics
+            }
+            with open(HISTORY_FILE, 'w') as f:
+                json.dump(data, f)
+            logger.info("Saved review history")
+        except Exception as e:
+            logger.error(f"Error saving history: {e}")
         
     @spaces.GPU
     def ensure_initialized(self):
@@ -60,14 +105,12 @@ class CodeReviewer:
                 login(token=HF_TOKEN, add_to_git_credential=False)
             
             logger.info("Loading tokenizer...")
-            # Initialize tokenizer with special tokens
             self.tokenizer = AutoTokenizer.from_pretrained(
                 MODEL_NAME,
                 token=HF_TOKEN,
                 trust_remote_code=True,
                 cache_dir=CACHE_DIR
             )
-            # Ensure special tokens are set
             special_tokens = {
                 'pad_token': '[PAD]',
                 'eos_token': '</s>',
@@ -87,13 +130,13 @@ class CodeReviewer:
                 cache_dir=CACHE_DIR,
                 token=HF_TOKEN
             )
-            # Resize embeddings for special tokens if needed
             if num_added > 0:
                 logger.info("Resizing model embeddings for special tokens")
                 self.model.resize_token_embeddings(len(self.tokenizer))
             
             self.device = next(self.model.parameters()).device
             logger.info(f"Model loaded successfully on {self.device}")
+            self._initialized = True
             return True
         except Exception as e:
             logger.error(f"Error initializing model: {e}")
@@ -117,14 +160,12 @@ Code:
     def review_code(self, code: str, language: str) -> str:
         """Perform code review using the model."""
         try:
-            # Ensure model is initialized
             if not self._initialized and not self.initialize_model():
                 return "Error: Model initialization failed. Please try again later."
             
             start_time = datetime.now()
             prompt = self.create_review_prompt(code, language)
             
-            # Tokenize with error handling
             try:
                 inputs = self.tokenizer(
                     prompt,
@@ -140,7 +181,6 @@ Code:
                 logger.error(f"Tokenization error: {token_error}")
                 return "Error: Failed to process input code. Please try again."
             
-            # Generate with error handling
             try:
                 with torch.no_grad():
                     outputs = self.model.generate(
@@ -158,7 +198,6 @@ Code:
                 logger.error(f"Generation error: {gen_error}")
                 return "Error: Failed to generate review. Please try again."
             
-            # Decode with error handling
             try:
                 response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
                 suggestions = response[len(prompt):].strip()
@@ -166,16 +205,14 @@ Code:
                 logger.error(f"Decoding error: {decode_error}")
                 return "Error: Failed to decode model output. Please try again."
             
-            # Create review and update metrics
             end_time = datetime.now()
             review = Review(code, language, suggestions)
             review.response_time = (end_time - start_time).total_seconds()
             self.review_history.append(review)
             
-            # Update metrics
             self.update_metrics(review)
+            self.save_history()  # Save after each review
             
-            # Clear GPU memory
             if self.device and self.device.type == "cuda":
                 del inputs, outputs
                 torch.cuda.empty_cache()
@@ -190,12 +227,10 @@ Code:
         """Update metrics with new review."""
         self.metrics['total_reviews'] += 1
         
-        # Update average response time
         total_time = self.metrics['avg_response_time'] * (self.metrics['total_reviews'] - 1)
         total_time += review.response_time
         self.metrics['avg_response_time'] = total_time / self.metrics['total_reviews']
         
-        # Update reviews today
         today = datetime.now().date()
         self.metrics['reviews_today'] = sum(
             1 for r in self.review_history 
@@ -212,7 +247,7 @@ Code:
                 'suggestions': r.suggestions,
                 'response_time': f"{r.response_time:.2f}s"
             }
-            for r in reversed(self.review_history[-10:])  # Last 10 reviews
+            for r in reversed(self.review_history[-10:])
         ]
 
     def get_metrics(self) -> Dict:
@@ -266,13 +301,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
                 label="Performance Metrics"
             )
     
-    # Set up event handlers
     @spaces.GPU
     def review_code_interface(code: str, language: str) -> str:
         if not code.strip():
             return "Please enter some code to review."
         try:
-            reviewer.ensure_initialized()  # Ensure model is initialized
+            reviewer.ensure_initialized()
             return reviewer.review_code(code, language)
         except Exception as e:
             logger.error(f"Interface error: {e}")
@@ -317,7 +351,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
         outputs=metrics_output
     )
     
-    # Add example inputs
     gr.Examples(
         examples=[
             ["""def add_numbers(a, b):
@@ -333,7 +366,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
         inputs=[code_input, language_input]
     )
 
-# Launch the app
 if __name__ == "__main__":
     iface.launch(
         server_name="0.0.0.0",