Update app.py

app.py

@@ -35,98 +35,95 @@ 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."
         
-        # 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
+        # ULTIMATE FIX: Check if CPU version exists, if not convert it
+        cpu_pickle_path = pickle_path.replace('.pkl', '_cpu.pkl')
         
-        # Also set other CUDA-disabling flags
-        os_module.environ['CUDA_LAUNCH_BLOCKING'] = '0'
-        
-        try:
-            # 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)
+        if not os.path.exists(cpu_pickle_path):
+            # Need to convert CUDA pickle to CPU pickle
+            try:
+                # Use torch.load with custom map_location that captures and remaps ALL devices
+                def smart_map_location(storage, location):
+                    # This function is called for EACH tensor storage during unpickling
+                    # It runs BEFORE the "CUDA device check", allowing us to remap
+                    return storage.cpu()
+                
+                # Load with our smart mapper
+                model_package = torch.load(pickle_path, map_location=smart_map_location)
+                
+                # Now save it as a CPU-only pickle for future loads
+                torch.save(model_package, cpu_pickle_path)
+                
+                return f"✅ Converted CUDA model to CPU! Loading from converted version...\n\nPlease wait, loading model..."
+                
+            except Exception as convert_error:
+                # Conversion failed, try direct load with aggressive remapping
+                import io
+                import pickle as pkl
                 
-                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:
-                # 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)
+                    # Read the pickle bytes
+                    with open(pickle_path, 'rb') as f:
+                        buffer = io.BytesIO(f.read())
+                    
+                    # Create custom unpickler with aggressive CPU forcing
+                    class AggressiveCPUUnpickler(pkl.Unpickler):
+                        def find_class(self, module, name):
+                            # Remap any CUDA storage to CPU storage
+                            if 'cuda' in name.lower():
+                                name = name.replace('cuda', '').replace('Cuda', '')
+                            return super().find_class(module, name)
                         
-                        # 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())}"
+                        def load_build(self):
+                            # Override to catch tensor builds
+                            stack = self.stack
+                            state = stack.pop()
+                            inst = stack[-1]
                             
-                            # 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
+                            # If this is a tensor, force to CPU
+                            if hasattr(inst, 'to'):
+                                try:
+                                    inst = inst.cpu()
+                                    stack[-1] = inst
+                                except:
+                                    pass
                             
-                            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:
-            # 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)
+                            if hasattr(inst, '__setstate__'):
+                                inst.__setstate__(state)
+                            else:
+                                for k, v in state.items():
+                                    setattr(inst, k, v)
+                        
+                        def persistent_load(self, pid):
+                            # Intercept ALL storage loads
+                            if isinstance(pid, tuple) and len(pid) >= 5:
+                                # Standard torch storage format
+                                tag, storage_type, key, location, size = pid[0], pid[1], pid[2], pid[3], pid[4]
+                                if tag == 'storage':
+                                    # Force location to CPU
+                                    return (tag, storage_type, key, 'cpu', size)
+                            return pid
+                    
+                    # Try to load with aggressive unpickler
+                    unpickler = AggressiveCPUUnpickler(buffer)
+                    model_package = unpickler.load()
+                    
+                    # Save as CPU version for next time
+                    torch.save(model_package, cpu_pickle_path)
+                    
+                except Exception as aggressive_error:
+                    return (f"❌ Failed to convert CUDA pickle to CPU.\n\n"
+                            f"Convert error: {str(convert_error)[:100]}\n"
+                            f"Aggressive error: {str(aggressive_error)[:100]}\n\n"
+                            f"Please re-save your model on a CPU machine:\n"
+                            f"```python\n"
+                            f"import torch\n"
+                            f"# Load your model\n"
+                            f"model = model.cpu()  # Move to CPU\n"
+                            f"torch.save({{'model': model, 'tokenizer': tokenizer, 'config': config}}, 'best_model.pkl')\n"
+                            f"```")
+        else:
+            # CPU version exists, load it directly
+            model_package = torch.load(cpu_pickle_path, map_location='cpu')
         
         # Success! Model loaded with one of the strategies above
         # Handle a few common package shapes.