Update app.py

app.py

@@ -35,42 +35,98 @@ def load_model_from_pickle(pickle_path="best_model.pkl"):
         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."
         
-        # NUCLEAR OPTION: Mock the entire CUDA module to prevent driver checks
-        import sys
-        import types
+        # METHOD 1: Set environment variable BEFORE any CUDA operations
+        # This prevents PyTorch from seeing ANY CUDA devices
+        import os as os_module
+        old_cuda_visible = os_module.environ.get('CUDA_VISIBLE_DEVICES', None)
+        os_module.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # Disable all CUDA devices
         
-        # Create a fake CUDA module that claims everything is available but does nothing
-        fake_cuda = types.ModuleType('cuda')
-        fake_cuda.is_available = lambda: True
-        fake_cuda.device_count = lambda: 1
-        fake_cuda.current_device = lambda: 0
-        fake_cuda.get_device_name = lambda x: "CPU (mocked as CUDA)"
-        fake_cuda.set_device = lambda x: None
-        fake_cuda.device = lambda x: types.SimpleNamespace(__enter__=lambda: None, __exit__=lambda *args: None)
-        fake_cuda.init = lambda: None
-        fake_cuda.is_initialized = lambda: True
-        fake_cuda._initialization_lock = types.SimpleNamespace(__enter__=lambda: None, __exit__=lambda *args: None)
-        
-        # Save original
-        original_cuda = torch.cuda
+        # Also set other CUDA-disabling flags
+        os_module.environ['CUDA_LAUNCH_BLOCKING'] = '0'
         
         try:
-            # Replace torch.cuda with our mock during loading
-            torch.cuda = fake_cuda
-            sys.modules['torch'].cuda = fake_cuda
-            
-            # Load with aggressive CPU mapping
+            # METHOD 2: Use pickle with restricted globals to prevent CUDA imports
+            import io
+            import pickle
+            
+            class CPUOnlyUnpickler(pickle.Unpickler):
+                def find_class(self, module, name):
+                    # Allow transformers and torch imports
+                    if module.startswith('transformers') or module.startswith('torch'):
+                        return super().find_class(module, name)
+                    return super().find_class(module, name)
+                
+                def persistent_load(self, pid):
+                    # Intercept torch storage and force CPU
+                    if isinstance(pid, tuple) and len(pid) > 0:
+                        if pid[0] == 'storage':
+                            # Format: ('storage', storage_type, key, location, size)
+                            storage_type = pid[1]
+                            key = pid[2]
+                            location = 'cpu'  # Force CPU location
+                            size = pid[4] if len(pid) > 4 else pid[3]
+                            # Rebuild with CPU location
+                            return super().persistent_load(('storage', storage_type, key, location, size))
+                    return super().persistent_load(pid)
+            
+            # Load using our custom unpickler
             with open(pickle_path, 'rb') as f:
-                model_package = torch.load(
-                    f, 
-                    map_location='cpu',
-                    weights_only=False
-                )
+                # First try: Custom unpickler with CUDA disabled
+                try:
+                    model_package = CPUOnlyUnpickler(f).load()
+                except Exception as e1:
+                    # Second try: Standard torch.load with map_location
+                    f.seek(0)  # Reset file pointer
+                    try:
+                        model_package = torch.load(
+                            f,
+                            map_location=torch.device('cpu'),
+                            weights_only=False
+                        )
+                    except Exception as e2:
+                        # Third try: Load with pickle directly and extract weights only
+                        f.seek(0)
+                        raw_package = pickle.load(f)
+                        
+                        # Try to extract model from various package formats
+                        if isinstance(raw_package, dict):
+                            if 'model' in raw_package:
+                                model_obj = raw_package['model']
+                            elif 'state_dict' in raw_package:
+                                return (f"❌ The pickle contains only state_dict. Please save the full model object.\n\n"
+                                       f"Use: torch.save({{'model': model, 'tokenizer': tokenizer, 'config': config}}, 'file.pkl')")
+                            else:
+                                return f"❌ Unknown pickle format. Keys found: {list(raw_package.keys())}"
+                            
+                            # Move model to CPU recursively
+                            def recursive_cpu(obj):
+                                if hasattr(obj, 'cpu'):
+                                    return obj.cpu()
+                                elif isinstance(obj, dict):
+                                    return {k: recursive_cpu(v) for k, v in obj.items()}
+                                elif isinstance(obj, (list, tuple)):
+                                    return type(obj)(recursive_cpu(item) for item in obj)
+                                return obj
+                            
+                            model_package = {
+                                'model': recursive_cpu(model_obj) if model_obj else None,
+                                'tokenizer': raw_package.get('tokenizer'),
+                                'config': raw_package.get('config', {})
+                            }
+                        else:
+                            # Package is the model itself
+                            model_package = {
+                                'model': recursive_cpu(raw_package),
+                                'tokenizer': None,
+                                'config': {}
+                            }
         
         finally:
-            # Always restore original CUDA module
-            torch.cuda = original_cuda
-            sys.modules['torch'].cuda = original_cuda
+            # Restore original CUDA_VISIBLE_DEVICES
+            if old_cuda_visible is not None:
+                os_module.environ['CUDA_VISIBLE_DEVICES'] = old_cuda_visible
+            else:
+                os_module.environ.pop('CUDA_VISIBLE_DEVICES', None)
         
         # Success! Model loaded with one of the strategies above
         # Handle a few common package shapes.