Update main.py

main.py

@@ -19,11 +19,13 @@ def load_model():
     return num_imputer, scaler, model
 
 def preprocess_input_data(input_data, num_imputer, scaler):
-    input_data_df = pd.DataFrame(input_data)
-    num_columns = [col for col in input_data_df.columns if input_data_df[col].dtype != 'object']
-    input_data_imputed_num = num_imputer.transform(input_data_df[num_columns])
-    input_scaled_df = pd.DataFrame(scaler.transform(input_data_imputed_num), columns=num_columns)
-    return input_scaled_df
+    input_data_df = pd.DataFrame([input_data], index=[0])  # Add index [0] to the DataFrame
+    input_data_scaled = scaler.transform(input_data_df)
+    input_data_scaled = pd.DataFrame(input_data_scaled, columns=input_data_df.columns)
+    input_data_imputed = num_imputer.transform(input_data_scaled)
+
+    return input_data_imputed
+
 
 @app.get("/")
 def read_root():