Update app.py

app.py

@@ -11,9 +11,13 @@ print(f"Using device: {device}")
 # Load the pickled model and tokenizer
 def load_model():
     try:
-        # Load from pickle files
+        # Load from pickle files with CPU mapping
+        print("Loading model from pickle files...")
         with open('summarization_model.pkl', 'rb') as f:
-            model = pickle.load(f)
+            # Use torch.load with map_location to handle CUDA->CPU conversion
+            import io
+            model_bytes = f.read()
+            model = torch.load(io.BytesIO(model_bytes), map_location='cpu')
 
         with open('tokenizer.pkl', 'rb') as f:
             tokenizer = pickle.load(f)
@@ -25,7 +29,7 @@ def load_model():
             "summarization",
             model=model,
             tokenizer=tokenizer,
-            device=0 if torch.cuda.is_available() else -1
+            device=-1  # Force CPU usage
         )
 
         return summarizer
@@ -37,13 +41,28 @@ def load_model():
             summarizer = pipeline(
                 "summarization",
                 model="./summarization_model",
-                device=0 if torch.cuda.is_available() else -1
+                device=-1  # Force CPU usage
             )
             print("✓ Model loaded from directory")
             return summarizer
         except Exception as e:
             print(f"Error loading model: {e}")
             return None
+    except Exception as e:
+        print(f"Error loading from pickle: {e}")
+        print("Trying fallback loading...")
+        try:
+            # Fallback: try loading from directory
+            summarizer = pipeline(
+                "summarization",
+                model="./summarization_model",
+                device=-1  # Force CPU usage
+            )
+            print("✓ Model loaded from directory (fallback)")
+            return summarizer
+        except Exception as e2:
+            print(f"Error loading model: {e2}")
+            return None
 
 # Load model at startup
 summarizer = load_model()