Improve code quality and suppress HF token warnings (#252)

Author: noahho

README.md

@@ -208,33 +208,17 @@ TabPFN automatically downloads model weights when first used. For offline usage:
      - macOS: `~/Library/Caches/tabpfn/`
      - Linux: `~/.cache/tabpfn/`
 
-**Quick Download Script**
+**Using the Provided Download Script**
 
-```python
-import requests
-from tabpfn.utils import _user_cache_dir
-import sys
-
-# Get default cache directory using TabPFN's internal function
-cache_dir = _user_cache_dir(platform=sys.platform)
-cache_dir.mkdir(parents=True, exist_ok=True)
-
-# Define models to download
-models = {
-    "tabpfn-v2-classifier.ckpt": "https://huggingface.co/Prior-Labs/TabPFN-v2-clf/resolve/main/tabpfn-v2-classifier.ckpt",
-    "tabpfn-v2-regressor.ckpt": "https://huggingface.co/Prior-Labs/TabPFN-v2-reg/resolve/main/tabpfn-v2-regressor.ckpt",
-}
-
-# Download each model
-for name, url in models.items():
-    path = cache_dir / name
-    print(f"Downloading {name} to {path}")
-    with open(path, "wb") as f:
-        f.write(requests.get(url).content)
+If you have the TabPFN repository, you can use the included script to download all models (including ensemble variants):
 
-print(f"Models downloaded to {cache_dir}")
+```bash
+# After installing TabPFN
+python scripts/download_all_models.py
 ```
 
+This script will download the main classifier and regressor models, as well as all ensemble variant models to your system's default cache directory.
+
 **Q: I'm getting a `pickle` error when loading the model. What should I do?**  
 A: Try the following:
 - Download the newest version of tabpfn `pip install tabpfn --upgrade`

examples/tabpfn_for_binary_classification.py

@@ -1,4 +1,9 @@
 #  Copyright (c) Prior Labs GmbH 2025.
+"""Example of using TabPFN for binary classification.
+
+This example demonstrates how to use TabPFNClassifier on a binary classification task
+using the breast cancer dataset from scikit-learn.
+"""
 
 from sklearn.datasets import load_breast_cancer
 from sklearn.metrics import accuracy_score, roc_auc_score

examples/tabpfn_for_multiclass_classification.py

@@ -1,4 +1,9 @@
 #  Copyright (c) Prior Labs GmbH 2025.
+"""Example of using TabPFN for multiclass classification.
+
+This example demonstrates how to use TabPFNClassifier on a multiclass classification task
+using the Iris dataset from scikit-learn.
+"""
 
 from sklearn.datasets import load_iris
 from sklearn.metrics import accuracy_score, roc_auc_score

examples/tabpfn_for_regression.py

@@ -1,4 +1,9 @@
 #  Copyright (c) Prior Labs GmbH 2025.
+"""Example of using TabPFN for regression.
+
+This example demonstrates how to use TabPFNRegressor on a regression task
+using the diabetes dataset from scikit-learn.
+"""
 
 from sklearn.datasets import load_diabetes
 from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

scripts/__init__.py

@@ -0,0 +1,3 @@
+"""Scripts package for TabPFN."""
+
+from __future__ import annotations

scripts/download_all_models.py

@@ -0,0 +1,27 @@
+"""Download all TabPFN model files for offline use."""
+
+from __future__ import annotations
+
+import logging
+import sys
+
+from tabpfn.model.loading import _user_cache_dir, download_all_models
+
+
+def main() -> None:
+    """Download all TabPFN models and save to cache directory."""
+    # Configure logging
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    logger = logging.getLogger(__name__)
+
+    # Get default cache directory using TabPFN's internal function
+    cache_dir = _user_cache_dir(platform=sys.platform, appname="tabpfn")
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    logger.info(f"Downloading all models to {cache_dir}")
+    download_all_models(cache_dir)
+    logger.info(f"All models downloaded to {cache_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/get_max_dependencies.py

@@ -1,13 +1,24 @@
+"""Generate requirements.txt with maximum allowed dependency versions."""
+
+from __future__ import annotations
+
 import re
+from pathlib import Path
+
 
 def main() -> None:
-    with open('pyproject.toml', 'r') as f:
+    """Extract maximum dependency versions and write to requirements.txt."""
+    with Path("pyproject.toml").open() as f:
         content = f.read()
 
     # Find dependencies section using regex
-    deps_match = re.search(r'dependencies\s*=\s*\[(.*?)\]', content, re.DOTALL)
+    deps_match = re.search(r"dependencies\s*=\s*\[(.*?)\]", content, re.DOTALL)
     if deps_match:
-        deps = [d.strip(' "\'') for d in deps_match.group(1).strip().split('\n') if d.strip()]
+        deps = [
+            d.strip(' "\'')
+            for d in deps_match.group(1).strip().split("\n")
+            if d.strip()
+        ]
         max_reqs = []
         for dep in deps:
             # Check for maximum version constraint
@@ -18,11 +29,12 @@ def main() -> None:
                 max_reqs.append(f"{package}<{max_ver}")
             else:
                 # If no max version, just use the package name
-                package = re.match(r'([^>=<\s]+)', dep).group(1)
+                package = re.match(r"([^>=<\s]+)", dep).group(1)
                 max_reqs.append(package)
 
-        with open('requirements.txt', 'w') as f:
-            f.write('\n'.join(max_reqs))
+        with Path("requirements.txt").open("w") as f:
+            f.write("\n".join(max_reqs))
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

scripts/get_min_dependencies.py

@@ -1,22 +1,34 @@
+"""Generate requirements.txt with minimum dependency versions."""
+
+from __future__ import annotations
+
 import re
+from pathlib import Path
+
 
 def main() -> None:
-    with open('pyproject.toml', 'r') as f:
+    """Extract minimum dependency versions and write to requirements.txt."""
+    with Path("pyproject.toml").open() as f:
         content = f.read()
 
     # Find dependencies section using regex
-    deps_match = re.search(r'dependencies\s*=\s*\[(.*?)\]', content, re.DOTALL)
+    deps_match = re.search(r"dependencies\s*=\s*\[(.*?)\]", content, re.DOTALL)
     if deps_match:
-        deps = [d.strip(' "\'') for d in deps_match.group(1).strip().split('\n') if d.strip()]
+        deps = [
+            d.strip(' "\'')
+            for d in deps_match.group(1).strip().split("\n")
+            if d.strip()
+        ]
         min_reqs = []
         for dep in deps:
             match = re.match(r'([^>=<\s]+)\s*>=\s*([^,\s"\']+)', dep)
             if match:
                 package, min_ver = match.groups()
                 min_reqs.append(f"{package}=={min_ver}")
 
-        with open('requirements.txt', 'w') as f:
-            f.write('\n'.join(min_reqs))
+        with Path("requirements.txt").open("w") as f:
+            f.write("\n".join(min_reqs))
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

src/tabpfn/base.py

@@ -25,10 +25,8 @@
     InferenceEngineCachePreprocessing,
     InferenceEngineOnDemand,
 )
-from tabpfn.utils import (
-    infer_fp16_inference_mode,
-    load_model_criterion_config,
-)
+from tabpfn.model.loading import load_model_criterion_config
+from tabpfn.utils import infer_fp16_inference_mode
 
 if TYPE_CHECKING:
     import numpy as np