toggle pretraining limits based on number of samples

Author: Anup Kumar

tools/tabpfn/main.py

@@ -84,6 +84,8 @@ def train_evaluate(args):
     """
     Train TabPFN and predict
     """
+    MAX_IGNORE_PRETRAINING_LIMITS_SAMPLES = 1000
+    SEED = 42
     # prepare train data
     tr_features, tr_labels = separate_features_labels(args["train_data"], args["train_header"])
     # prepare test data
@@ -94,7 +96,10 @@ def train_evaluate(args):
         te_labels = []
     s_time = time.time()
     if args["selected_task"] == "Classification":
-        classifier = TabPFNClassifier(random_state=42, model_path=args["model_path"], ignore_pretraining_limits=True)
+        if tr_features.shape[0] <= MAX_IGNORE_PRETRAINING_LIMITS_SAMPLES:
+            classifier = TabPFNClassifier(random_state=SEED, model_path=args["model_path"])
+        else:
+            classifier = TabPFNClassifier(random_state=SEED, model_path=args["model_path"], ignore_pretraining_limits=True)
         classifier.fit(tr_features, tr_labels)
         y_eval = classifier.predict(te_features)
         pred_probas_test = classifier.predict_proba(te_features)
@@ -105,7 +110,10 @@ def train_evaluate(args):
             "output_predicted_data", sep="\t", index=None
         )
     else:
-        regressor = TabPFNRegressor(random_state=42, model_path=args["model_path"], ignore_pretraining_limits=True)
+        if tr_features.shape[0] <= MAX_IGNORE_PRETRAINING_LIMITS_SAMPLES:
+            regressor = TabPFNRegressor(random_state=SEED, model_path=args["model_path"])
+        else:
+            regressor = TabPFNRegressor(random_state=SEED, model_path=args["model_path"], ignore_pretraining_limits=True)
         regressor.fit(tr_features, tr_labels)
         y_eval = regressor.predict(te_features)
         if len(te_labels) > 0: