Move expense functionality to common (#94)

* Move expense functionality to common

* Ruff clean up

* Fix latency offset

* Update ref

Author: liam-sbhoo

tabpfn_client/mock_prediction.py

@@ -7,7 +7,13 @@
 import warnings
 import logging
 import functools
+
 from tabpfn_client.client import ServiceClient
+from tabpfn_client.tabpfn_common_utils.expense_estimation import estimate_duration
+
+
+COST_ESTIMATION_LATENCY_OFFSET = 1.0
+
 
 # For seamlessly switching between a mock mode for simulating prediction
 # costs and real prediction, use thread-local variables to keep track of the
@@ -49,53 +55,6 @@ def increment_mock_time(seconds: float):
     set_mock_time(get_mock_time() + seconds)
 
 
-def estimate_duration(
-    num_rows: int,
-    num_features: int,
-    task: Literal["classification", "regression"],
-    tabpfn_config: dict = {},
-) -> float:
-    """
-    Estimates the duration of a prediction task.
-    """
-    # Logic comes from _estimate_model_usage in base.py of the TabPFN codebase.
-    CONSTANT_COMPUTE_OVERHEAD = 8000
-    NUM_SAMPLES_FACTOR = 4
-    NUM_SAMPLES_PLUS_FEATURES = 6.5
-    CELLS_FACTOR = 0.25
-    CELLS_SQUARED_FACTOR = 1.3e-7
-    EMBEDDING_SIZE = 192
-    NUM_HEADS = 6
-    NUM_LAYERS = 12
-    FEATURES_PER_GROUP = 2
-    GPU_FACTOR = 1e-11
-    LATENCY_OFFSET = 1.0
-
-    n_estimators = tabpfn_config.get(
-        "n_estimators", 4 if task == "classification" else 8
-    )
-
-    num_samples = num_rows
-    num_feature_groups = int(np.ceil(num_features / FEATURES_PER_GROUP))
-
-    num_cells = (num_feature_groups + 1) * num_samples
-    compute_cost = (EMBEDDING_SIZE**2) * NUM_HEADS * NUM_LAYERS
-
-    base_duration = (
-        n_estimators
-        * compute_cost
-        * (
-            CONSTANT_COMPUTE_OVERHEAD
-            + num_samples * NUM_SAMPLES_FACTOR
-            + (num_samples + num_feature_groups) * NUM_SAMPLES_PLUS_FEATURES
-            + num_cells * CELLS_FACTOR
-            + num_cells**2 * CELLS_SQUARED_FACTOR
-        )
-    )
-
-    return round(base_duration * GPU_FACTOR + LATENCY_OFFSET, 3)
-
-
 def mock_predict(
     X_test,
     task: Literal["classification", "regression"],
@@ -116,7 +75,11 @@ def mock_predict(
         )
 
     duration = estimate_duration(
-        X_train.shape[0] + X_test.shape[0], X_test.shape[1], task, config
+        num_rows=X_train.shape[0] + X_test.shape[0],
+        num_features=X_test.shape[1],
+        task=task,
+        tabpfn_config=config,
+        latency_offset=COST_ESTIMATION_LATENCY_OFFSET,  # To slightly overestimate (safer)
     )
     increment_mock_time(duration)
 

tabpfn_client/tabpfn_common_utils

@@ -9,6 +9,7 @@
     get_mock_cost,
     estimate_duration,
     is_mock_mode,
+    COST_ESTIMATION_LATENCY_OFFSET,
 )
 from tabpfn_client.estimator import TabPFNClassifier, TabPFNRegressor
 
@@ -50,10 +51,11 @@ def test_mock_mode_behavior(self):
 
                 # Verify time increased by the estimated duration
                 expected_duration = estimate_duration(
-                    self.X_train.shape[0] + self.X_test.shape[0],
-                    self.X_test.shape[1],
-                    "classification",
-                    self.config,
+                    num_rows=self.X_train.shape[0] + self.X_test.shape[0],
+                    num_features=self.X_test.shape[1],
+                    task="classification",
+                    tabpfn_config=self.config,
+                    latency_offset=COST_ESTIMATION_LATENCY_OFFSET,
                 )
                 self.assertAlmostEqual(
                     get_mock_time() - initial_time, expected_duration