Add pandas ai as cleaning method

experiment.py

@@ -34,7 +34,7 @@ def one_search_experiment(dataset, error_type, train_file, model, seed, n_jobs=1
     result = train_and_evaluate(X_train, y_train, X_test_list, y_test_list, test_files, model, n_jobs=n_jobs, seed=train_seed, hyperparams=hyperparams)
     return result
 
-def one_split_experiment(dataset, n_retrain=5, seed=1, n_jobs=1, nosave=True, error_type=None):
+def one_split_experiment(dataset, n_retrain=1, seed=1, n_jobs=1, nosave=True, error_type=None):
     """Run experiments on one dataset for one split.
 
     Args:

schema/clean_method.py

@@ -1,6 +1,9 @@
 # define the domain of cleaning method
+import base64
 import numpy as np
 import pandas as pd
+from pandasai.llm import OpenAI
+from pandasai import SmartDataframe
 from sklearn.neighbors import LocalOutlierFactor
 from sklearn.ensemble import IsolationForest
 from sklearn import preprocessing
@@ -10,6 +13,7 @@
 import utils
 import os
 
+
 class MVCleaner(object):
     def __init__(self, method='delete', **kwargs):
         self.method = method
@@ -68,6 +72,57 @@ def clean(self, dirty_train, dirty_test):
         clean_test, indicator_test = self.clean_df(dirty_test)
         return clean_train, indicator_train, clean_test, indicator_test
 
+class PandasAICleaner:
+    def __init__(self, open_ai_api_key="YOUR_API_TOKEN"):
+        self.open_ai_api_key = "dummy" 
+        self.llm = OpenAI(api_token=self.open_ai_api_key)
+
+    def fit(self, dataset, dirty_train):
+        # Instantiate an LLM using the provided API key
+        pass
+
+    def detect(self, df):
+        return df.isnull()
+
+    def clean_df(self, df):    
+        mv_mat = self.detect(df)
+        self.smart_df = SmartDataframe(df, config={"llm": self.llm})
+        df_clean = pd.DataFrame(list(self.smart_df.impute_missing_values()))
+        return df_clean, mv_mat
+
+    def clean(self, dirty_train, dirty_test):
+        clean_train, indicator_train = self.clean_df(dirty_train) #detect and repair
+        clean_test, indicator_test = self.clean_df(dirty_test)
+        return clean_train, indicator_train, clean_test, indicator_test
+
+
+class MVHoloCleaner(object):
+    def __init__(self):
+        self.tag = "impute_holoclean"
+
+    def detect(self, df):
+        return df.isnull()
+
+    def fit(self, dataset, df):
+        clean_raw_path = utils.get_dir(dataset, 'raw', 'Holoclean_mv_clean.csv')
+        clean_raw = pd.read_csv(clean_raw_path)
+
+        index_train_path = utils.get_dir(dataset, 'raw', 'idx_train.csv')
+        index_test_path = utils.get_dir(dataset, 'raw', 'idx_test.csv')
+        index_train = pd.read_csv(index_train_path).values.reshape(-1)
+        index_test = pd.read_csv(index_test_path).values.reshape(-1)
+
+        self.clean_train = clean_raw.iloc[index_train, :]
+        self.clean_test = clean_raw.iloc[index_test, :]
+
+    def clean(self, dirty_train, dirty_test):
+        indicator_train = self.detect(dirty_train)
+        indicator_test = self.detect(dirty_test)
+
+        clean_train = self.clean_train
+        clean_test =self.clean_test
+        return clean_train, indicator_train, clean_test, indicator_test
+
 class DuplicatesCleaner(object):
     def __init__(self):
         super(DuplicatesCleaner, self).__init__()
@@ -390,32 +445,6 @@ def clean(self, dirty_train, dirty_test):
         clean_test, indicator_test = self.clean_df(dirty_test)
         return clean_train, indicator_train, clean_test, indicator_test
 
-class MVHoloCleaner(object):
-    def __init__(self):
-        self.tag = "impute_holoclean"
-
-    def detect(self, df):
-        return df.isnull()
-
-    def fit(self, dataset, df):
-        clean_raw_path = utils.get_dir(dataset, 'raw', 'Holoclean_mv_clean.csv')
-        clean_raw = pd.read_csv(clean_raw_path)
-
-        index_train_path = utils.get_dir(dataset, 'raw', 'idx_train.csv')
-        index_test_path = utils.get_dir(dataset, 'raw', 'idx_test.csv')
-        index_train = pd.read_csv(index_train_path).values.reshape(-1)
-        index_test = pd.read_csv(index_test_path).values.reshape(-1)
-
-        self.clean_train = clean_raw.iloc[index_train, :]
-        self.clean_test = clean_raw.iloc[index_test, :]
-
-    def clean(self, dirty_train, dirty_test):
-        indicator_train = self.detect(dirty_train)
-        indicator_test = self.detect(dirty_test)
-
-        clean_train = self.clean_train
-        clean_test =self.clean_test
-        return clean_train, indicator_train, clean_test, indicator_test
 
 class MVHumanCleaner(object):
     def __init__(self):

schema/dataset.py

@@ -18,7 +18,7 @@
     "ml_task": "classification"
 }
 
-Airbnb = {
+Airbnb1 = {
     "data_dir": "Airbnb",
     "error_types": ["duplicates", "outliers", "missing_values"],
     "label": 'Rating',
@@ -27,6 +27,15 @@
     'key_columns': ['latitude', 'longitude'],
 }
 
+Airbnb = {
+    "data_dir": "Airbnb",
+    "error_types": ["missing_values"],
+    "label": 'Rating',
+    "categorical_variables": ['Rating'],
+    "ml_task": "classification",
+    'key_columns': ['latitude', 'longitude'],
+}
+
 Titanic = {
     "data_dir": "Titanic",
     "error_types": ["missing_values"],

schema/error_type.py

@@ -4,14 +4,16 @@
 # details of each error type
 missing_values = {
     "name": "missing_values",
-    "clean_methods": {"delete": MVCleaner("delete"),
-                      "impute_holoclean": MVHoloCleaner(),
-                      "impute_mean_mode": MVCleaner("impute", num="mean", cat="mode"),
-                      "impute_mean_dummy": MVCleaner("impute", num="mean", cat="dummy"),
-                      "impute_median_mode": MVCleaner("impute", num="median", cat="mode"),
-                      "impute_median_dummy": MVCleaner("impute", num="median", cat="dummy"),
-                      "impute_mode_mode": MVCleaner("impute", num="mode", cat="mode"),
-                      "impute_mode_dummy": MVCleaner("impute", num="mode", cat="dummy"),
+    "clean_methods": {
+                      "delete": MVCleaner("delete"),
+                      "impute_pandas_ai": PandasAICleaner()
+                    #   "impute_holoclean": MVHoloCleaner(),
+                    #   "impute_mean_mode": MVCleaner("impute", num="mean", cat="mode"),
+                    #   "impute_mean_dummy": MVCleaner("impute", num="mean", cat="dummy"),
+                    #   "impute_median_mode": MVCleaner("impute", num="median", cat="mode"),
+                    #   "impute_median_dummy": MVCleaner("impute", num="median", cat="dummy"),
+                    #   "impute_mode_mode": MVCleaner("impute", num="mode", cat="mode"),
+                    #   "impute_mode_dummy": MVCleaner("impute", num="mode", cat="dummy"),
                     }
 }
 

schema/model.py

@@ -93,4 +93,4 @@
 }
 
 # model domain
-models = [logistic_reg, knn_clf, dt_clf, adaboost_clf, random_forest_clf, gaussian_nb, xgb_clf]
+models = [logistic_reg] #, knn_clf, dt_clf, adaboost_clf, random_forest_clf, gaussian_nb, xgb_clf]