Update app.py

app.py

@@ -104,6 +104,44 @@ Then upload 'best_model_cpu.pkl' to this Space and rename it to 'best_model.pkl'
             return ("❌ No model object found inside the pickle. Please ensure the pickle contains a dict with keys "
                     "'model', 'tokenizer', and 'config' (or the model object itself).")
         
+        # Fix tokenizer compatibility issues
+        if loaded_tokenizer is not None:
+            try:
+                # Ensure tokenizer has required attributes for generation
+                if not hasattr(loaded_tokenizer, 'pad_token_id') or loaded_tokenizer.pad_token_id is None:
+                    loaded_tokenizer.pad_token_id = loaded_tokenizer.eos_token_id
+                
+                # Fix missing _unk_token attribute (common in older tokenizers)
+                if not hasattr(loaded_tokenizer, '_unk_token'):
+                    if hasattr(loaded_tokenizer, 'unk_token'):
+                        loaded_tokenizer._unk_token = loaded_tokenizer.unk_token
+                    else:
+                        loaded_tokenizer._unk_token = '<unk>'
+                
+                # Ensure other critical attributes exist
+                if not hasattr(loaded_tokenizer, '_bos_token'):
+                    loaded_tokenizer._bos_token = getattr(loaded_tokenizer, 'bos_token', '<s>')
+                if not hasattr(loaded_tokenizer, '_eos_token'):
+                    loaded_tokenizer._eos_token = getattr(loaded_tokenizer, 'eos_token', '</s>')
+                
+                # Test tokenizer basic functionality
+                test_encode = loaded_tokenizer("test", return_tensors='pt')
+                test_decode = loaded_tokenizer.decode(test_encode['input_ids'][0])
+                
+            except Exception as tokenizer_error:
+                # Tokenizer is broken, try to recreate it
+                try:
+                    from transformers import GPT2Tokenizer
+                    print(f"⚠️ Loaded tokenizer has issues ({tokenizer_error}), recreating from GPT-2...")
+                    loaded_tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+                    
+                    # Ensure pad token is set
+                    if loaded_tokenizer.pad_token_id is None:
+                        loaded_tokenizer.pad_token_id = loaded_tokenizer.eos_token_id
+                        
+                except Exception as recreate_error:
+                    return f"❌ Tokenizer error: {tokenizer_error}\nRecreation failed: {recreate_error}\n\nPlease ensure the tokenizer is compatible with current transformers version."
+        
         # Set model to evaluation mode and move to appropriate device
         try:
             loaded_model.eval()
@@ -207,23 +245,38 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
         # Format input
         prompt = f"<PSEUDO> {pseudo_code.strip()} <SEP> <CODE>"
         
-        # Tokenize
+        # Tokenize with error handling
         device = next(loaded_model.parameters()).device
-        inputs = loaded_tokenizer(prompt, return_tensors='pt').to(device)
+        try:
+            inputs = loaded_tokenizer(prompt, return_tensors='pt').to(device)
+        except Exception as tokenize_error:
+            # Try to fix tokenizer on the fly
+            try:
+                from transformers import GPT2Tokenizer
+                print("Fixing tokenizer compatibility...")
+                loaded_tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+                if loaded_tokenizer.pad_token_id is None:
+                    loaded_tokenizer.pad_token_id = loaded_tokenizer.eos_token_id
+                inputs = loaded_tokenizer(prompt, return_tensors='pt').to(device)
+            except Exception as fix_error:
+                return f"❌ Tokenization failed: {tokenize_error}\nFix attempt failed: {fix_error}", "", "", ""
         
         # Generate (ensure type safety for parameters)
         with torch.no_grad():
-            outputs = loaded_model.generate(
-                **inputs,
-                max_length=int(max_length),
-                temperature=float(temperature),
-                top_k=int(top_k),
-                top_p=float(top_p),
-                do_sample=True,
-                num_return_sequences=int(num_sequences),
-                pad_token_id=loaded_tokenizer.pad_token_id,
-                eos_token_id=loaded_tokenizer.eos_token_id,
-            )
+            try:
+                outputs = loaded_model.generate(
+                    **inputs,
+                    max_length=int(max_length),
+                    temperature=float(temperature),
+                    top_k=int(top_k),
+                    top_p=float(top_p),
+                    do_sample=True,
+                    num_return_sequences=int(num_sequences),
+                    pad_token_id=loaded_tokenizer.pad_token_id,
+                    eos_token_id=loaded_tokenizer.eos_token_id,
+                )
+            except Exception as generation_error:
+                return f"❌ Generation failed: {str(generation_error)}", "", "", ""
         
         generation_time = time.time() - start_time