Improve model initialization with better state management

app.py

@@ -44,9 +44,15 @@ class CodeReviewer:
             'avg_response_time': 0.0,
             'reviews_today': 0
         }
-        self.initialize_model()
+        self._initialized = False
         
     @spaces.GPU
+    def ensure_initialized(self):
+        """Ensure model is initialized."""
+        if not self._initialized:
+            self.initialize_model()
+            self._initialized = True
+        
     def initialize_model(self):
         """Initialize the model and tokenizer."""
         try:
@@ -67,7 +73,8 @@ class CodeReviewer:
                 'eos_token': '</s>',
                 'bos_token': '<s>'
             }
-            self.tokenizer.add_special_tokens(special_tokens)
+            num_added = self.tokenizer.add_special_tokens(special_tokens)
+            logger.info(f"Added {num_added} special tokens")
             logger.info("Tokenizer loaded successfully")
             
             logger.info("Loading model...")
@@ -80,13 +87,18 @@ class CodeReviewer:
                 cache_dir=CACHE_DIR,
                 token=HF_TOKEN
             )
-            # Resize embeddings for special tokens
-            self.model.resize_token_embeddings(len(self.tokenizer))
+            # 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}")
+            return True
         except Exception as e:
             logger.error(f"Error initializing model: {e}")
-            raise
+            self._initialized = False
+            return False
 
     def create_review_prompt(self, code: str, language: str) -> str:
         """Create a structured prompt for code review."""
@@ -104,10 +116,11 @@ Code:
     @spaces.GPU
     def review_code(self, code: str, language: str) -> str:
         """Perform code review using the model."""
-        if not self.tokenizer or not self.model:
-            return "Error: Model not properly initialized. Please try again later."
-            
         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)
             
@@ -163,7 +176,7 @@ Code:
             self.update_metrics(review)
             
             # Clear GPU memory
-            if self.device.type == "cuda":
+            if self.device and self.device.type == "cuda":
                 del inputs, outputs
                 torch.cuda.empty_cache()
             
@@ -208,7 +221,7 @@ Code:
             'Total Reviews': self.metrics['total_reviews'],
             'Average Response Time': f"{self.metrics['avg_response_time']:.2f}s",
             'Reviews Today': self.metrics['reviews_today'],
-            'Device': str(self.device)
+            'Device': str(self.device) if self.device else "Not initialized"
         }
 
 # Initialize reviewer
@@ -259,6 +272,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
         if not code.strip():
             return "Please enter some code to review."
         try:
+            reviewer.ensure_initialized()  # Ensure model is initialized
             return reviewer.review_code(code, language)
         except Exception as e:
             logger.error(f"Interface error: {e}")