class EstimatorAdapter: def __init__(self, estimator): self.estimator = estimator def predict(self, batch): return self.estimator.predict(batch) def predict_proba(self, batch): return self.estimator.predict_proba(batch)