fbcode/openzl/dev/tools/py/zstrong_ml.py

Reviewed By: terrelln

Differential Revision: D114197886

Author: DevmateMrMediaCqsF405 Bot

tools/py/zstrong_ml.py

@@ -7,7 +7,35 @@
 
 import pandas as pd
 import xgboost as xgb
-from zstrong_json.ml import *
+from zstrong_json.ml import (
+    BaseFeatureGenerator as BaseFeatureGenerator,
+    DeltaIntFeatureGenerator as DeltaIntFeatureGenerator,
+    FeatureGenerator as FeatureGenerator,
+    GBTModel as GBTModel,
+    IntFeatureGenerator as IntFeatureGenerator,
+    MLModel as MLModel,
+    MLSelector as MLSelector,
+    MLTrainingSample as MLTrainingSample,
+    MLTrainingSelector as MLTrainingSelector,
+    samples_from_json as samples_from_json,
+    samples_to_json as samples_to_json,
+    TokenizeIntFeatureGenerator as TokenizeIntFeatureGenerator,
+)
+
+_ZSTRONG_JSON_ML_REEXPORTS = (
+    BaseFeatureGenerator,
+    DeltaIntFeatureGenerator,
+    FeatureGenerator,
+    GBTModel,
+    IntFeatureGenerator,
+    MLModel,
+    MLSelector,
+    MLTrainingSample,
+    MLTrainingSelector,
+    samples_from_json,
+    samples_to_json,
+    TokenizeIntFeatureGenerator,
+)
 
 
 class CoreModel(t.NamedTuple):