update2

example/train.py

@@ -173,6 +173,9 @@ def main():
     std = stats["std"]
     transforms = stats.get("transform", {})
     raw_std = stats.get("raw_std", std)
+    quantile_probs = stats.get("quantile_probs")
+    quantile_values = stats.get("quantile_values")
+    use_quantile = bool(config.get("use_quantile_transform", False))
 
     vocab = load_json(config["vocab_path"])["vocab"]
     vocab_sizes = [len(vocab[c]) for c in disc_cols]
@@ -253,6 +256,9 @@ def main():
                     max_batches=int(config["max_batches"]),
                     return_file_id=False,
                     transforms=transforms,
+                    quantile_probs=quantile_probs,
+                    quantile_values=quantile_values,
+                    use_quantile=use_quantile,
                     shuffle_buffer=int(config.get("shuffle_buffer", 0)),
                 )
             ):
@@ -284,6 +290,9 @@ def main():
                 max_batches=int(config["max_batches"]),
                 return_file_id=use_condition,
                 transforms=transforms,
+                quantile_probs=quantile_probs,
+                quantile_values=quantile_values,
+                use_quantile=use_quantile,
                 shuffle_buffer=int(config.get("shuffle_buffer", 0)),
             )
         ):