Update app.py

app.py

@@ -34,20 +34,50 @@ 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
+        
+        # Strategy: Use a custom unpickler that forces all CUDA tensors to CPU during deserialization
+        # This prevents the "Attempting to deserialize object on a CUDA device" error entirely
+        class CPUUnpickler(pickle.Unpickler):
+            def find_class(self, module, name):
+                # Remap torch CUDA classes to CPU equivalents
+                if module == 'torch.storage' and name.endswith('Storage'):
+                    # Force all storage types to CPU (e.g., cuda:0 -> cpu)
+                    return super().find_class('torch', name.replace('Cuda', '').replace('cuda', ''))
+                return super().find_class(module, name)
+            
+            def persistent_load(self, pid):
+                # Handle torch storage tuples: ('storage', storage_type, key, location, size)
+                if isinstance(pid, tuple) and len(pid) > 0 and pid[0] == 'storage':
+                    storage_type, key, location, size = pid[1], pid[2], pid[3], pid[4]
+                    # Force location to CPU regardless of original device
+                    if 'cuda' in str(location).lower():
+                        location = 'cpu'
+                    # Rebuild the storage descriptor with CPU location
+                    return ('storage', storage_type, key, location, size)
+                return pid
+        
         try:
-            if torch.cuda.is_available():
-                model_package = torch.load(pickle_path)
-            else:
+            # Attempt 1: Custom CPU unpickler (most robust for CUDA->CPU)
+            with open(pickle_path, 'rb') as f:
+                model_package = CPUUnpickler(f).load()
+        except Exception as e1:
+            try:
+                # Attempt 2: torch.load with map_location (standard approach)
                 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
+            except Exception as e2:
+                try:
+                    # Attempt 3: torch.load with weights_only=False (for older torch versions)
+                    model_package = torch.load(pickle_path, map_location='cpu', weights_only=False)
+                except Exception as e3:
+                    # All strategies failed - return combined error info
+                    return (f"❌ All loading strategies failed:\n\n"
+                            f"1. Custom CPU unpickler: {str(e1)[:100]}\n"
+                            f"2. torch.load(map_location): {str(e2)[:100]}\n"
+                            f"3. torch.load(weights_only=False): {str(e3)[:100]}\n\n"
+                            f"The pickle may be corrupted or incompatible with this PyTorch version.\n"
+                            f"Try re-saving with: model.to('cpu'); torch.save({{'model': model, ...}}, 'file.pkl')")
         
+        # Success! Model loaded with one of the strategies above
         # Handle a few common package shapes.
         if isinstance(model_package, dict):
             loaded_model = model_package.get('model', None)
@@ -100,18 +130,9 @@ def load_model_from_pickle(pickle_path="best_model.pkl"):
         return config_info
         
     except Exception as e:
-        # Specific hint for CUDA->CPU deserialization issues
+        # Final catch-all for any unexpected errors
         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."
+        return f"❌ Unexpected error loading model: {err}\n\nPlease ensure best_model.pkl is properly uploaded and compatible with this environment."
 
 def calculate_bleu_score(reference, hypothesis):
     """Calculate BLEU score between reference and generated code"""