ruff

Author: SamuelGabriel

quick_test.py

@@ -1,10 +1,9 @@
 import logging
-import numpy as np
 
 from sklearn.datasets import load_breast_cancer, load_diabetes
 from sklearn.model_selection import train_test_split
 
-from tabpfn_client import estimator, UserDataClient, init
+from tabpfn_client import UserDataClient, init
 from tabpfn_client.estimator import TabPFNClassifier, TabPFNRegressor
 
 logging.basicConfig(level=logging.DEBUG)
@@ -44,4 +43,4 @@
     print("predicting reg")
     print(tabpfn.predict(X_test))
 
-    print(UserDataClient().get_data_summary())
+    print(UserDataClient().get_data_summary())

tabpfn_client/client.py

@@ -100,7 +100,13 @@ def upload_train_set(self, X, y) -> str:
         train_set_uid = response.json()["train_set_uid"]
         return train_set_uid
 
-    def predict(self, train_set_uid: str, x_test, task: Literal["classification", "regression"], tabpfn_config: dict | None = None) -> dict[str, np.ndarray]:
+    def predict(
+        self,
+        train_set_uid: str,
+        x_test,
+        task: Literal["classification", "regression"],
+        tabpfn_config: dict | None = None,
+    ) -> dict[str, np.ndarray]:
         """
         Predict the class labels for the provided data (test set).
 
@@ -140,7 +146,7 @@ def predict(self, train_set_uid: str, x_test, task: Literal["classification", "r
 
         if not isinstance(result, dict):
             result = {"probas": result}
-        
+
         for k in result:
             result[k] = np.array(result[k])
 

tabpfn_client/config.py

@@ -73,4 +73,4 @@ def reset():
         g_tabpfn_config.user_auth_handler.reset_cache()
 
     # remove cache dir
-    shutil.rmtree(CACHE_DIR, ignore_errors=True)
+    shutil.rmtree(CACHE_DIR, ignore_errors=True)

tabpfn_client/estimator.py

@@ -10,6 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 @dataclass(eq=True, frozen=True)
 class PreprocessorConfig:
     """
@@ -108,6 +109,7 @@ def to_dict(self):
             for k, v in asdict(self).items()
         }
 
+
 class TabPFNClassifier(BaseEstimator, ClassifierMixin):
     def __init__(
         self,
@@ -128,7 +130,9 @@ def __init__(
         feature_shift_decoder: str = "shuffle",
         normalize_with_test: bool = False,
         average_logits: bool = False,
-        optimize_metric: Literal["auroc", "roc", "auroc_ovo", "balanced_acc", "acc", "log_loss", None] = "roc",
+        optimize_metric: Literal[
+            "auroc", "roc", "auroc_ovo", "balanced_acc", "acc", "log_loss", None
+        ] = "roc",
         transformer_predict_kwargs: Optional[dict] = None,
         multiclass_decoder="shuffle",
         softmax_temperature: Optional[float] = -0.1,
@@ -205,7 +209,9 @@ def predict(self, X):
 
     def predict_proba(self, X):
         check_is_fitted(self)
-        return config.g_tabpfn_config.inference_handler.predict(X, task="classification", config=self.get_params())["probas"]
+        return config.g_tabpfn_config.inference_handler.predict(
+            X, task="classification", config=self.get_params()
+        )["probas"]
 
 
 class TabPFNRegressor(BaseEstimator, RegressorMixin):
@@ -225,7 +231,9 @@ def __init__(
         feature_shift_decoder: str = "shuffle",
         normalize_with_test: bool = False,
         average_logits: bool = False,
-        optimize_metric: Literal["mse", "rmse", "mae", "r2", "mean", "median", "mode", "exact_match", None] = "rmse",
+        optimize_metric: Literal[
+            "mse", "rmse", "mae", "r2", "mean", "median", "mode", "exact_match", None
+        ] = "rmse",
         transformer_predict_kwargs: Optional[Dict] = None,
         softmax_temperature: Optional[float] = -0.1,
         use_poly_features=False,
@@ -324,7 +332,7 @@ def fit(self, X, y):
                 "Only server mode is supported at the moment for tabpfn_classifier.init(use_server=False)"
             )
         return self
-    
+
     def predict(self, X):
         full_prediction_dict = self.predict_full(X)
         if self.optimize_metric in ("mse", "rmse", "r2", "mean", None):
@@ -335,7 +343,9 @@ def predict(self, X):
             return full_prediction_dict["mode"]
         else:
             raise ValueError(f"Optimize metric {self.optimize_metric} not supported")
-    
+
     def predict_full(self, X):
         check_is_fitted(self)
-        return config.g_tabpfn_config.inference_handler.predict(X, task="regression", config=self.get_params())
+        return config.g_tabpfn_config.inference_handler.predict(
+            X, task="regression", config=self.get_params()
+        )

tabpfn_client/service_wrapper.py

@@ -180,7 +180,10 @@ def fit(self, X, y) -> None:
 
         self.last_train_set_uid = self.service_client.upload_train_set(X, y)
 
-    def predict(self, X, task:Literal["classification","regression"], config=None):
+    def predict(self, X, task: Literal["classification", "regression"], config=None):
         return self.service_client.predict(
-            train_set_uid=self.last_train_set_uid, x_test=X, tabpfn_config=config, task=task,
+            train_set_uid=self.last_train_set_uid,
+            x_test=X,
+            tabpfn_config=config,
+            task=task,
         )

tabpfn_client/tests/integration/test_tabpfn_classifier.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from tabpfn_client import init, reset
-from tabpfn_client import estimator, TabPFNClassifier
+from tabpfn_client import TabPFNClassifier
 from tabpfn_client.tests.mock_tabpfn_server import with_mock_server
 from tabpfn_client.service_wrapper import UserAuthenticationClient
 from tabpfn_client.client import ServiceClient

tabpfn_client/tests/unit/test_client.py

@@ -179,7 +179,9 @@ def test_predict_with_valid_train_set_and_test_set(self, mock_server):
         )
 
         pred = self.client.predict(
-            train_set_uid=dummy_json["train_set_uid"], x_test=self.X_test, task="classification"
+            train_set_uid=dummy_json["train_set_uid"],
+            x_test=self.X_test,
+            task="classification",
         )
         self.assertTrue(np.array_equal(pred["probas"], dummy_result["classification"]))