Fix tokenizer initialization and improve error handling

app.py

@@ -54,12 +54,21 @@ 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>',
+                'bos_token': '<s>'
+            }
+            self.tokenizer.add_special_tokens(special_tokens)
+            logger.info("Tokenizer loaded successfully")
             
             logger.info("Loading model...")
             self.model = AutoModelForCausalLM.from_pretrained(
@@ -71,6 +80,8 @@ class CodeReviewer:
                 cache_dir=CACHE_DIR,
                 token=HF_TOKEN
             )
+            # Resize 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}")
         except Exception as e:
@@ -93,6 +104,9 @@ 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:
             start_time = datetime.now()
             prompt = self.create_review_prompt(code, language)
@@ -105,7 +119,10 @@ Code:
                     truncation=True,
                     max_length=512,
                     padding=True
-                ).to(self.device)
+                )
+                if inputs is None:
+                    raise ValueError("Failed to tokenize input")
+                inputs = inputs.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."
@@ -120,7 +137,9 @@ Code:
                         temperature=0.7,
                         top_p=0.95,
                         num_beams=1,
-                        early_stopping=True
+                        early_stopping=True,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        eos_token_id=self.tokenizer.eos_token_id
                     )
             except Exception as gen_error:
                 logger.error(f"Generation error: {gen_error}")