Update utils/data_cleaning.py

utils/data_cleaning.py

@@ -1,31 +1,40 @@
-import pandas as pd
-import numpy as np
-
-# Function to handle missing values
-def handle_missing_values(df, method="drop"):
-    if method == "drop":
-        df = df.dropna()
-    elif method == "fill":
-        for col in df.columns:
-            if df[col].dtype in ['float64', 'int64']:
-                df[col].fillna(df[col].mean(), inplace=True)
-            else:
-                df[col].fillna(df[col].mode()[0], inplace=True)
-    return df
-
-# Function to remove outliers using the IQR method
-def remove_outliers_iqr(df, column):
-    Q1 = df[column].quantile(0.25)
-    Q3 = df[column].quantile(0.75)
-    IQR = Q3 - Q1
-    lower_bound = Q1 - 1.5 * IQR
-    upper_bound = Q3 + 1.5 * IQR
-    return df[(df[column] >= lower_bound) & (df[column] <= upper_bound)]
-
-# Function to cap extreme values (5th and 95th percentiles)
-def cap_extreme_values(df):
-    for col in df.select_dtypes(include=[np.number]).columns:
-        lower_limit = df[col].quantile(0.05)
-        upper_limit = df[col].quantile(0.95)
-        df[col] = np.clip(df[col], lower_limit, upper_limit)
-    return df
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import LabelEncoder
+
+def handle_missing_values(df: pd.DataFrame):
+    """Handle missing values in the dataframe."""
+    fill_method = "Fill with mean/median"  # Can be dynamic (from user input)
+    if fill_method == "Drop rows":
+        df = df.dropna()
+    elif fill_method == "Fill with mean/median":
+        for col in df.columns:
+            if df[col].dtype in ['float64', 'int64']:
+                df[col].fillna(df[col].mean(), inplace=True)
+            else:
+                df[col].fillna(df[col].mode()[0], inplace=True)
+    return df
+
+def remove_outliers_iqr(df: pd.DataFrame):
+    """Remove outliers using the IQR method."""
+    Q1 = df.quantile(0.25)
+    Q3 = df.quantile(0.75)
+    IQR = Q3 - Q1
+    df = df[~((df < (Q1 - 1.5 * IQR)) | (df > (Q3 + 1.5 * IQR))).any(axis=1)]
+    return df
+
+def cap_extreme_values(df: pd.DataFrame):
+    """Cap extreme values using the 5th and 95th percentiles."""
+    for col in df.select_dtypes(include=[np.number]).columns:
+        lower_limit = df[col].quantile(0.05)
+        upper_limit = df[col].quantile(0.95)
+        df[col] = np.clip(df[col], lower_limit, upper_limit)
+    return df
+
+def encode_categorical_data(df: pd.DataFrame):
+    """Encode categorical columns to numeric."""
+    label_encoder = LabelEncoder()
+    for col in df.columns:
+        if df[col].dtype == 'object' or len(df[col].unique()) <= 10:
+            df[col] = label_encoder.fit_transform(df[col])
+    return df