most frequent imputer function

most_frequent_imputer.py

def ImputeWithMostFrequent(train_df, test_df, 
                           cols=["country", "province", 
                                 "region_1", "variety"]):
  # function to impute country, province, region_1, variety
  # columns with the most_frequent value of each feature
  # most_frequent imputer is fitted on train dataset
  # transformation done on the train and test set

  train_df = pd.DataFrame(train_df[cols])
  test_df = pd.DataFrame(test_df[cols]) 
    
  most_frequent_imputer = SimpleImputer(strategy="most_frequent")

  most_frequent_imputer.fit(train_df)
    
  imputed_train_set = most_frequent_imputer.transform(train_df)
  imputed_train_df = pd.DataFrame(imputed_train_set, columns=train_df.columns)

  imputed_test_set = most_frequent_imputer.transform(test_df)
  imputed_test_df = pd.DataFrame(imputed_test_set, columns=test_df.columns)

  return imputed_train_df, imputed_test_df