Update model_wrapper.py

model_wrapper.py

@@ -1,38 +1,155 @@
 import os
 import threading
 
+from transformers import (
+    AutoTokenizer, 
+    AutoModel, 
+    AutoModelForSeq2SeqLM, 
+    AutoModelForCausalLM,
+    AutoConfig
+)
+import torch
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 class CodeDebuggerWrapper:
-    """
-    Simple wrapper that loads the same HF model and exposes debug(code: str) -> str
-    This is used by app.py (Gradio).
-    """
-    def __init__(self, model_name: str = "Girinath11/aiml_code_debug_model"):
+    def __init__(self, model_name="Girinath11/aiml_code_debug_model"):
         self.model_name = model_name
-        self._lock = threading.Lock()
-        self.tokenizer = None
         self.model = None
-        self.max_new_tokens = int(os.environ.get("MAX_NEW_TOKENS", "256"))
+        self.tokenizer = None
+        self.model_type = None
         self._ensure_model()
-
+    
     def _ensure_model(self):
-        # allow skipping in environments where you don't want to download weights
-        skip = os.environ.get("SKIP_MODEL_LOAD", "0") == "1"
-        if skip:
-            print("SKIP_MODEL_LOAD=1 -> not loading model.")
-            return
-
-        if self.model is None or self.tokenizer is None:
-            from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-            with self._lock:
-                if self.model is None or self.tokenizer is None:
-                    print(f"Loading model {self.model_name} ...")
-                    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-                    self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
-                    print("Model loaded.")
-
+        """Load model and tokenizer with fallback strategies."""
+        logger.info(f"Loading model {self.model_name} ...")
+        
+        try:
+            # Load tokenizer first
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            logger.info("✅ Tokenizer loaded successfully")
+            
+            # Add pad token if it doesn't exist
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Check model configuration
+            config = AutoConfig.from_pretrained(self.model_name)
+            logger.info(f"Model config type: {type(config).__name__}")
+            
+            # Try different model loading strategies
+            loading_strategies = [
+                ("AutoModel with trust_remote_code", lambda: AutoModel.from_pretrained(
+                    self.model_name, trust_remote_code=True
+                )),
+                ("AutoModelForCausalLM with trust_remote_code", lambda: AutoModelForCausalLM.from_pretrained(
+                    self.model_name, trust_remote_code=True
+                )),
+                ("AutoModelForSeq2SeqLM with trust_remote_code", lambda: AutoModelForSeq2SeqLM.from_pretrained(
+                    self.model_name, trust_remote_code=True
+                )),
+                ("AutoModel without trust_remote_code", lambda: AutoModel.from_pretrained(
+                    self.model_name
+                )),
+                ("AutoModelForCausalLM without trust_remote_code", lambda: AutoModelForCausalLM.from_pretrained(
+                    self.model_name
+                )),
+            ]
+            
+            for strategy_name, strategy_func in loading_strategies:
+                try:
+                    logger.info(f"Trying: {strategy_name}")
+                    self.model = strategy_func()
+                    self.model_type = type(self.model).__name__
+                    logger.info(f"✅ Successfully loaded model with {strategy_name}")
+                    logger.info(f"Model type: {self.model_type}")
+                    break
+                except Exception as e:
+                    logger.warning(f"❌ {strategy_name} failed: {str(e)[:100]}...")
+                    continue
+            
+            if self.model is None:
+                raise RuntimeError("❌ Failed to load model with any strategy")
+            
+            # Set to evaluation mode
+            self.model.eval()
+            logger.info("✅ Model set to evaluation mode")
+            
+        except Exception as e:
+            logger.error(f"❌ Critical error in model loading: {e}")
+            raise
+    
     def debug(self, code: str) -> str:
-        if self.model is None or self.tokenizer is None:
-            return "Model not loaded. Set SKIP_MODEL_LOAD=0 and ensure HF token is available if model is private."
-        inputs = self.tokenizer(code, return_tensors="pt", padding=True, truncation=True)
-        outputs = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens)
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        """Debug the provided code."""
+        if not code or not code.strip():
+            return "❌ Please provide some code to debug."
+        
+        try:
+            # Prepare input with a clear prompt
+            prompt = f"### Task: Debug and fix the following Python code\n\n### Input Code:\n{code}\n\n### Fixed Code:"
+            
+            # Tokenize input
+            inputs = self.tokenizer(
+                prompt, 
+                return_tensors="pt", 
+                max_length=512, 
+                truncation=True,
+                padding=True
+            )
+            
+            logger.info(f"Input tokenized. Input IDs shape: {inputs['input_ids'].shape}")
+            
+            # Generate response based on model type
+            with torch.no_grad():
+                if hasattr(self.model, 'generate'):
+                    logger.info("Using generate method")
+                    
+                    # Generation parameters
+                    generation_kwargs = {
+                        **inputs,
+                        'max_new_tokens': 256,
+                        'num_beams': 3,
+                        'early_stopping': True,
+                        'do_sample': True,
+                        'temperature': 0.7,
+                        'pad_token_id': self.tokenizer.pad_token_id,
+                    }
+                    
+                    # Add eos_token_id if available
+                    if hasattr(self.tokenizer, 'eos_token_id') and self.tokenizer.eos_token_id is not None:
+                        generation_kwargs['eos_token_id'] = self.tokenizer.eos_token_id
+                    
+                    try:
+                        outputs = self.model.generate(**generation_kwargs)
+                        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+                        
+                        # Clean up the response
+                        if prompt in response:
+                            response = response.replace(prompt, "").strip()
+                        
+                        # If response is empty or too short, provide a fallback
+                        if not response or len(response.strip()) < 10:
+                            response = f"The model processed your code but didn't generate a clear fix. Original code:\n\n{code}\n\n💡 Tip: Check for common issues like indentation, syntax errors, or missing imports."
+                        
+                    except Exception as gen_error:
+                        logger.error(f"Generation error: {gen_error}")
+                        response = f"❌ Error during code generation: {str(gen_error)}\n\nOriginal code:\n{code}"
+                
+                else:
+                    # For encoder-only models
+                    logger.info("Model doesn't support generation. Using encoder-only approach.")
+                    outputs = self.model(**inputs)
+                    
+                    # This is a simplified approach - you might need to add a classification head
+                    # or use a different strategy based on your specific model
+                    response = f"⚠️  This model type ({self.model_type}) doesn't support text generation directly.\n\nOriginal code:\n{code}\n\n💡 Consider using a generative model (T5, GPT, BART) for code debugging tasks."
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"Error in debug method: {e}")
+            return f"❌ Error during debugging: {str(e)}\n\nOriginal code:\n{code}"
+