Add robust error handling and improve model loading

app.py

@@ -7,6 +7,10 @@ import logging
 from datetime import datetime
 import json
 from typing import List, Dict
+import warnings
+
+# Filter CUDA warnings
+warnings.filterwarnings('ignore', category=UserWarning, message='Can\'t initialize NVML')
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -16,6 +20,10 @@ logger = logging.getLogger(__name__)
 HF_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
 MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
 
+# Cache directory for model
+CACHE_DIR = "/home/user/.cache/huggingface"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
 class Review:
     def __init__(self, code: str, language: str, suggestions: str):
         self.code = code
@@ -47,18 +55,36 @@ class CodeReviewer:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 MODEL_NAME,
                 token=HF_TOKEN,
-                trust_remote_code=True
+                trust_remote_code=True,
+                cache_dir=CACHE_DIR
             )
             
             logger.info("Loading model...")
-            self.model = AutoModelForCausalLM.from_pretrained(
-                MODEL_NAME,
-                token=HF_TOKEN,
-                device_map="auto",
-                torch_dtype=torch.float16,
-                trust_remote_code=True,
-                low_cpu_mem_usage=True
-            )
+            # Initialize model with specific configuration
+            model_kwargs = {
+                "device_map": "auto",
+                "torch_dtype": torch.float16,
+                "trust_remote_code": True,
+                "low_cpu_mem_usage": True,
+                "cache_dir": CACHE_DIR,
+                "token": HF_TOKEN
+            }
+            
+            # Load model with error handling
+            try:
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    MODEL_NAME,
+                    **model_kwargs
+                )
+            except Exception as model_error:
+                logger.error(f"Error loading model: {model_error}")
+                # Try loading with safetensors
+                model_kwargs["use_safetensors"] = True
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    MODEL_NAME,
+                    **model_kwargs
+                )
+            
             logger.info(f"Model loaded successfully on {self.device}")
         except Exception as e:
             logger.error(f"Error initializing model: {e}")
@@ -83,27 +109,42 @@ Code:
             start_time = datetime.now()
             prompt = self.create_review_prompt(code, language)
             
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                truncation=True,
-                max_length=512,
-                padding=True
-            ).to(self.device)
+            # Tokenize with error handling
+            try:
+                inputs = self.tokenizer(
+                    prompt,
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=512,
+                    padding=True
+                ).to(self.device)
+            except Exception as token_error:
+                logger.error(f"Tokenization error: {token_error}")
+                return "Error: Failed to process input code. Please try again."
             
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=512,
-                    do_sample=True,
-                    temperature=0.7,
-                    top_p=0.95,
-                    num_beams=1,
-                    early_stopping=True
-                )
+            # Generate with error handling
+            try:
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=512,
+                        do_sample=True,
+                        temperature=0.7,
+                        top_p=0.95,
+                        num_beams=1,
+                        early_stopping=True
+                    )
+            except Exception as gen_error:
+                logger.error(f"Generation error: {gen_error}")
+                return "Error: Failed to generate review. Please try again."
             
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            suggestions = response[len(prompt):].strip()
+            # Decode with error handling
+            try:
+                response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+                suggestions = response[len(prompt):].strip()
+            except Exception as decode_error:
+                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()
@@ -212,24 +253,33 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
         try:
             return reviewer.review_code(code, language)
         except Exception as e:
+            logger.error(f"Interface error: {e}")
             return f"Error: {str(e)}"
     
     def get_history_interface() -> str:
-        history = reviewer.get_history()
-        if not history:
-            return "No reviews yet."
-        result = ""
-        for review in history:
-            result += f"Time: {review['timestamp']}\n"
-            result += f"Language: {review['language']}\n"
-            result += f"Response Time: {review['response_time']}\n"
-            result += "Code:\n```\n" + review['code'] + "\n```\n"
-            result += "Suggestions:\n" + review['suggestions'] + "\n"
-            result += "-" * 80 + "\n\n"
-        return result
+        try:
+            history = reviewer.get_history()
+            if not history:
+                return "No reviews yet."
+            result = ""
+            for review in history:
+                result += f"Time: {review['timestamp']}\n"
+                result += f"Language: {review['language']}\n"
+                result += f"Response Time: {review['response_time']}\n"
+                result += "Code:\n```\n" + review['code'] + "\n```\n"
+                result += "Suggestions:\n" + review['suggestions'] + "\n"
+                result += "-" * 80 + "\n\n"
+            return result
+        except Exception as e:
+            logger.error(f"History error: {e}")
+            return "Error retrieving history"
     
     def get_metrics_interface() -> Dict:
-        return reviewer.get_metrics()
+        try:
+            return reviewer.get_metrics()
+        except Exception as e:
+            logger.error(f"Metrics error: {e}")
+            return {"error": str(e)}
     
     submit_btn.click(
         review_code_interface,