code: def setup_predictor(self, df, target):
      mask = np.random.rand(len(df)) < 0.8
      train = df[mask]
      test = df[~mask]
  
      predictor = predictor_from_problem(
          df,
          ProblemDefinition.from_dict(
              {"target": target, "time_aim": 20, "unbias_target": False}
          ),
      )
  
      predictor.learn(train)
  
      if hasattr(predictor, "ensemble"):
          for i, mixer in enumerate(predictor.ensemble.mixers):
              predictor.ensemble.best_index = i
              predictions = predictor.predict(test)
              assert "prediction" in predictions.columns
              assert "confidence" in predictions.columns
  
              predictions = predictor.predict(
                  test, args={"predict_proba": True}
              )
  
              for label in df[target].unique():
                  assert f"__mdb_proba_{label}" in predictions.columns
      return predictor
  
file_path: lightwood/tests/integration/basic/test_categorical.py
imports: import unittest
  import numpy
  import pandas
  from sklearn.metrics import balanced_accuracy_score
  from lightwood.api.types import ProblemDefinition
  from lightwood.api.high_level import predictor_from_problem
label_desc: Use image node_id 1 for calling the TestBasic obj's underlying member method code with example usage: obj.setup_predictor(df, target) and returns: predictor