Update app.py

app.py

@@ -77,19 +77,24 @@ def load_data():
     # Replace prognosis values with numerical categories
     df.replace({'prognosis': disease_dict}, inplace=True)
 
-    # Check unique values in the prognosis column to capture any unmapped entries
+    # Unique values for debugging
     print("Unique values in prognosis after mapping:", df['prognosis'].unique())
 
-    # Ensure prognosis is purely numerical after mapping
-    df['prognosis'] = df['prognosis'].astype(int)  # Convert to integer if necessary
-    
+    # Before converting to integer, check if all entries are mapped correctly
+    if df['prognosis'].dtype == 'object':  # Check if any mapping error occurred
+        raise ValueError(f"Prognosis contains unmapped values: {df['prognosis'].unique()}")
+
+    # Convert to integer
+    df['prognosis'] = df['prognosis'].astype(int)
+
     # Inference doesn't require fixing as copy=True defaults
     df = df.infer_objects()
 
     tr.replace({'prognosis': disease_dict}, inplace=True)
     
     # Ensure it is also numerical
-    print("Unique values in prognosis for testing data after mapping:", tr['prognosis'].unique())
+    if tr['prognosis'].dtype == 'object':
+        raise ValueError(f"Testing data prognosis contains unmapped values: {tr['prognosis'].unique()}")
     
     tr['prognosis'] = tr['prognosis'].astype(int)  # Convert to integer if necessary
     tr = tr.infer_objects()
@@ -105,7 +110,7 @@ y_test = tr['prognosis']
 
 # Encode the target variable with LabelEncoder if not already numerical
 le = LabelEncoder()
-y_encoded = le.fit_transform(y)  # Needs to be string labels, assuming df['prognosis'] has no numerical labels
+y_encoded = le.fit_transform(y)  # Fits and transforms the labels
 
 def train_models():
     models = {