Update app.py

app.py

@@ -34,21 +34,51 @@ def load_model_from_pickle(pickle_path="best_model.pkl"):
         # Check if file exists
         if not os.path.exists(pickle_path):
             return f"❌ Model file not found: {pickle_path}\n\nPlease ensure best_model.pkl is uploaded to the HuggingFace Space."
+        # Try to load using torch.load which supports map_location for CPU-only machines
+        try:
+            if torch.cuda.is_available():
+                model_package = torch.load(pickle_path)
+            else:
+                model_package = torch.load(pickle_path, map_location=torch.device('cpu'))
+        except RuntimeError as rte:
+            # Common error when a GPU-saved object is loaded on CPU-only machine
+            if 'Attempting to deserialize object on a CUDA device' in str(rte):
+                # Retry mapping to CPU
+                model_package = torch.load(pickle_path, map_location=torch.device('cpu'))
+            else:
+                raise
         
-        # Load pickle file
-        with open(pickle_path, 'rb') as f:
-            model_package = pickle.load(f)
-        
-        loaded_model = model_package['model']
-        loaded_tokenizer = model_package['tokenizer']
-        loaded_config = model_package['config']
-        
-        # Set model to evaluation mode
-        loaded_model.eval()
+        # Handle a few common package shapes.
+        if isinstance(model_package, dict):
+            loaded_model = model_package.get('model', None)
+            loaded_tokenizer = model_package.get('tokenizer', None)
+            loaded_config = model_package.get('config', {}) or {}
+        else:
+            # Unknown package format: assume the object itself is the model
+            loaded_model = model_package
+            loaded_tokenizer = None
+            loaded_config = {}
+
+        # If user saved a state_dict instead of a model object, provide guidance
+        if isinstance(loaded_model, dict) and 'state_dict' in loaded_model:
+            # the file contains something like {'state_dict': ...}
+            return ("❌ The pickle appears to contain a state_dict rather than a full model object. "
+                    "This app expects a pickled model object (model instance).\n"
+                    "If you only have a state_dict, re-create the model architecture and load the state_dict before pickling, "
+                    "or provide a pickled model object saved with torch.save(model, path).")
+
+        if loaded_model is None:
+            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).")
         
-        # Move to appropriate device
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        loaded_model = loaded_model.to(device)
+        # Set model to evaluation mode and move to appropriate device
+        try:
+            loaded_model.eval()
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            loaded_model = loaded_model.to(device)
+        except Exception as e:
+            return (f"❌ Error preparing model for inference: {str(e)}\n\n"
+                    "This can happen if the saved object is not a proper torch.nn.Module or if tensors couldn't be mapped to the current device.")
         
         config_info = f"""✅ Model loaded successfully!
 
@@ -70,7 +100,18 @@ def load_model_from_pickle(pickle_path="best_model.pkl"):
         return config_info
         
     except Exception as e:
-        return f"❌ Error loading model: {str(e)}\n\nPlease ensure best_model.pkl is properly uploaded to the Space."
+        # Specific hint for CUDA->CPU deserialization issues
+        err = str(e)
+        if 'Attempting to deserialize object on a CUDA device' in err:
+            return ("❌ Error loading model: The file was saved for a GPU device but this runtime has no CUDA available. "
+                    "The loader attempted to remap tensors to CPU, but loading still failed.\n\n"
+                    "Try re-saving the model on a CPU by running:\n"
+                    "```python\n"
+                    "model.to('cpu')\n"
+                    "torch.save({'model': model, 'tokenizer': tokenizer, 'config': config}, 'best_model.pkl')\n"
+                    "```\n\n"
+                    "Or upload a CPU-compatible `best_model.pkl` to the Space and retry.")
+        return f"❌ Error loading model: {err}\n\nPlease ensure best_model.pkl is properly uploaded to the Space and is CPU-compatible."
 
 def calculate_bleu_score(reference, hypothesis):
     """Calculate BLEU score between reference and generated code"""