smallaxe 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. smallaxe-0.1.0/LICENSE +21 -0
  2. smallaxe-0.1.0/PKG-INFO +204 -0
  3. smallaxe-0.1.0/README.md +156 -0
  4. smallaxe-0.1.0/pyproject.toml +102 -0
  5. smallaxe-0.1.0/setup.cfg +4 -0
  6. smallaxe-0.1.0/smallaxe/__init__.py +157 -0
  7. smallaxe-0.1.0/smallaxe/_config.py +37 -0
  8. smallaxe-0.1.0/smallaxe/auto/__init__.py +1 -0
  9. smallaxe-0.1.0/smallaxe/datasets/__init__.py +13 -0
  10. smallaxe-0.1.0/smallaxe/datasets/_data.py +240 -0
  11. smallaxe-0.1.0/smallaxe/exceptions/__init__.py +120 -0
  12. smallaxe-0.1.0/smallaxe/metrics/__init__.py +35 -0
  13. smallaxe-0.1.0/smallaxe/metrics/classification.py +241 -0
  14. smallaxe-0.1.0/smallaxe/metrics/regression.py +301 -0
  15. smallaxe-0.1.0/smallaxe/pipeline/__init__.py +5 -0
  16. smallaxe-0.1.0/smallaxe/pipeline/pipeline.py +691 -0
  17. smallaxe-0.1.0/smallaxe/preprocessing/__init__.py +11 -0
  18. smallaxe-0.1.0/smallaxe/preprocessing/encoder.py +410 -0
  19. smallaxe-0.1.0/smallaxe/preprocessing/imputer.py +327 -0
  20. smallaxe-0.1.0/smallaxe/preprocessing/scaler.py +285 -0
  21. smallaxe-0.1.0/smallaxe/search/__init__.py +1 -0
  22. smallaxe-0.1.0/smallaxe/training/__init__.py +16 -0
  23. smallaxe-0.1.0/smallaxe/training/base.py +764 -0
  24. smallaxe-0.1.0/smallaxe/training/classifiers.py +127 -0
  25. smallaxe-0.1.0/smallaxe/training/mixins/__init__.py +15 -0
  26. smallaxe-0.1.0/smallaxe/training/mixins/metadata_mixin.py +158 -0
  27. smallaxe-0.1.0/smallaxe/training/mixins/param_mixin.py +151 -0
  28. smallaxe-0.1.0/smallaxe/training/mixins/persistence_mixin.py +164 -0
  29. smallaxe-0.1.0/smallaxe/training/mixins/spark_model_mixin.py +255 -0
  30. smallaxe-0.1.0/smallaxe/training/mixins/validation_mixin.py +228 -0
  31. smallaxe-0.1.0/smallaxe/training/random_forest.py +198 -0
  32. smallaxe-0.1.0/smallaxe/training/regressors.py +125 -0
  33. smallaxe-0.1.0/smallaxe/viz/__init__.py +1 -0
  34. smallaxe-0.1.0/smallaxe.egg-info/PKG-INFO +204 -0
  35. smallaxe-0.1.0/smallaxe.egg-info/SOURCES.txt +50 -0
  36. smallaxe-0.1.0/smallaxe.egg-info/dependency_links.txt +1 -0
  37. smallaxe-0.1.0/smallaxe.egg-info/requires.txt +25 -0
  38. smallaxe-0.1.0/smallaxe.egg-info/top_level.txt +1 -0
  39. smallaxe-0.1.0/tests/test_config.py +240 -0
  40. smallaxe-0.1.0/tests/test_datasets.py +274 -0
  41. smallaxe-0.1.0/tests/test_encoder.py +569 -0
  42. smallaxe-0.1.0/tests/test_exceptions.py +219 -0
  43. smallaxe-0.1.0/tests/test_factories.py +374 -0
  44. smallaxe-0.1.0/tests/test_imputer.py +298 -0
  45. smallaxe-0.1.0/tests/test_metrics_classification.py +524 -0
  46. smallaxe-0.1.0/tests/test_metrics_regression.py +390 -0
  47. smallaxe-0.1.0/tests/test_mixins.py +610 -0
  48. smallaxe-0.1.0/tests/test_pipeline.py +666 -0
  49. smallaxe-0.1.0/tests/test_random_forest.py +1008 -0
  50. smallaxe-0.1.0/tests/test_scaler.py +359 -0
  51. smallaxe-0.1.0/tests/test_smoke.py +62 -0
  52. smallaxe-0.1.0/tests/test_training_base.py +530 -0
smallaxe-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 Henok Yemam
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,204 @@
1
+ Metadata-Version: 2.4
2
+ Name: smallaxe
3
+ Version: 0.1.0
4
+ Summary: A PySpark MLOps library for simplified model training and optimization
5
+ Author: Henok Yemam
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/henokyemam/smallaxe
8
+ Project-URL: Repository, https://github.com/henokyemam/smallaxe
9
+ Project-URL: Issues, https://github.com/henokyemam/smallaxe/issues
10
+ Keywords: pyspark,mlops,machine-learning,spark,xgboost,lightgbm,catboost
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Operating System :: OS Independent
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3.8
18
+ Classifier: Programming Language :: Python :: 3.9
19
+ Classifier: Programming Language :: Python :: 3.10
20
+ Classifier: Programming Language :: Python :: 3.11
21
+ Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
23
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
24
+ Requires-Python: <3.13,>=3.8
25
+ Description-Content-Type: text/markdown
26
+ License-File: LICENSE
27
+ Requires-Dist: pyspark<4.0,>=3.3
28
+ Requires-Dist: hyperopt>=0.2.7
29
+ Requires-Dist: plotly>=5.0.0
30
+ Provides-Extra: xgboost
31
+ Requires-Dist: xgboost>=1.7.0; extra == "xgboost"
32
+ Provides-Extra: lightgbm
33
+ Requires-Dist: lightgbm>=3.3.0; extra == "lightgbm"
34
+ Requires-Dist: synapse-ml-lightgbm>=1.0.0; extra == "lightgbm"
35
+ Provides-Extra: catboost
36
+ Requires-Dist: catboost>=1.1.0; extra == "catboost"
37
+ Provides-Extra: all
38
+ Requires-Dist: smallaxe[xgboost]; extra == "all"
39
+ Requires-Dist: smallaxe[lightgbm]; extra == "all"
40
+ Requires-Dist: smallaxe[catboost]; extra == "all"
41
+ Provides-Extra: dev
42
+ Requires-Dist: pytest>=7.0.0; extra == "dev"
43
+ Requires-Dist: pytest-spark>=0.6.0; extra == "dev"
44
+ Requires-Dist: black>=23.0.0; extra == "dev"
45
+ Requires-Dist: ruff>=0.1.0; extra == "dev"
46
+ Requires-Dist: mypy>=1.0.0; extra == "dev"
47
+ Dynamic: license-file
48
+
49
+ # smallaxe
50
+
51
+ [![CI](https://github.com/henokyemam/smallaxe/actions/workflows/ci.yml/badge.svg)](https://github.com/henokyemam/smallaxe/actions/workflows/ci.yml)
52
+
53
+ A PySpark MLOps library that simplifies model training, evaluation, and optimization for PySpark DataFrames.
54
+
55
+ ## Why smallaxe?
56
+
57
+ PySpark MLlib has a steep learning curve and verbose API. **smallaxe** provides a clean, scikit-learn-like interface for common ML workflows while leveraging the distributed power of Spark.
58
+
59
+ ## Features
60
+
61
+ - **Simple API** - Train models with familiar `fit()`/`predict()` patterns
62
+ - **Multiple Algorithms** - XGBoost, LightGBM, CatBoost, and Random Forest
63
+ - **Preprocessing Pipeline** - Imputer, Scaler, Encoder with chainable pipelines
64
+ - **Hyperparameter Optimization** - Built-in hyperopt integration with early stopping
65
+ - **Automated Training** - Train all algorithms and compare with one call
66
+ - **Visualization** - Plotly-based charts for model evaluation
67
+ - **Cross-Validation** - Train/test split and k-fold with stratified sampling
68
+
69
+ ## Installation
70
+
71
+ ```bash
72
+ pip install smallaxe
73
+ ```
74
+
75
+ Install with optional algorithm dependencies:
76
+
77
+ ```bash
78
+ pip install smallaxe[xgboost] # XGBoost support
79
+ pip install smallaxe[lightgbm] # LightGBM support
80
+ pip install smallaxe[catboost] # CatBoost support
81
+ pip install smallaxe[all] # All algorithms
82
+ ```
83
+
84
+ ## Quick Start
85
+
86
+ ```python
87
+ from smallaxe.training import Regressors
88
+ from smallaxe.datasets import load_sample_regression
89
+
90
+ # Load sample data
91
+ df = load_sample_regression(spark)
92
+
93
+ # Train a model
94
+ model = Regressors.random_forest()
95
+ model.fit(df, label_col='price', exclude_cols=['id'])
96
+
97
+ # Make predictions
98
+ predictions = model.predict(df)
99
+ ```
100
+
101
+ ## Usage Examples
102
+
103
+ ### Training with Cross-Validation
104
+
105
+ ```python
106
+ from smallaxe.training import Classifiers
107
+
108
+ model = Classifiers.xgboost(task='binary')
109
+ model.fit(
110
+ df,
111
+ label_col='churn',
112
+ validation='kfold',
113
+ n_folds=5,
114
+ stratified=True
115
+ )
116
+
117
+ print(model.validation_scores)
118
+ ```
119
+
120
+ ### Preprocessing Pipeline
121
+
122
+ ```python
123
+ from smallaxe.pipeline import Pipeline
124
+ from smallaxe.preprocessing import Imputer, Scaler, Encoder
125
+ from smallaxe.training import Regressors
126
+
127
+ pipeline = Pipeline([
128
+ ('imputer', Imputer(numerical_strategy='median')),
129
+ ('scaler', Scaler(method='standard')),
130
+ ('encoder', Encoder(method='onehot')),
131
+ ('model', Regressors.xgboost())
132
+ ])
133
+
134
+ pipeline.fit(
135
+ df,
136
+ label_col='target',
137
+ numerical_cols=['age', 'income'],
138
+ categorical_cols=['city', 'category']
139
+ )
140
+
141
+ predictions = pipeline.predict(new_df)
142
+ ```
143
+
144
+ ### Hyperparameter Optimization
145
+
146
+ ```python
147
+ from smallaxe.search import optimize
148
+ from hyperopt import hp
149
+
150
+ param_grid = {
151
+ 'max_depth': hp.choice('max_depth', [3, 5, 7, 10]),
152
+ 'learning_rate': hp.uniform('learning_rate', 0.01, 0.3)
153
+ }
154
+
155
+ best_model = optimize.run(
156
+ model=Regressors.xgboost(),
157
+ dataframe=df,
158
+ label_col='target',
159
+ param_grid=param_grid,
160
+ metric='rmse',
161
+ max_evals=50
162
+ )
163
+
164
+ print(best_model.best_params)
165
+ ```
166
+
167
+ ### Automated Training
168
+
169
+ ```python
170
+ from smallaxe.auto import AutomatedTraining
171
+
172
+ auto = AutomatedTraining(model_type='classification', metrics=['f1_score', 'auc_roc'])
173
+ auto.fit(
174
+ df,
175
+ label_col='churn',
176
+ numerical_cols=['tenure', 'monthly_charges'],
177
+ categorical_cols=['contract'],
178
+ n_folds=5
179
+ )
180
+
181
+ # Compare all models
182
+ auto.metrics.show()
183
+
184
+ # Use best model
185
+ predictions = auto.predict(new_df)
186
+ ```
187
+
188
+ ## Supported Algorithms
189
+
190
+ | Algorithm | Regressor | Classifier | Dependencies |
191
+ |-----------|-----------|------------|--------------|
192
+ | Random Forest | Yes | Yes | None (native PySpark) |
193
+ | XGBoost | Yes | Yes | `smallaxe[xgboost]` |
194
+ | LightGBM | Yes | Yes | `smallaxe[lightgbm]` |
195
+ | CatBoost | Yes | Yes | `smallaxe[catboost]` |
196
+
197
+ ## Requirements
198
+
199
+ - Python 3.8 - 3.12
200
+ - PySpark 3.3+
201
+
202
+ ## License
203
+
204
+ MIT License
@@ -0,0 +1,156 @@
1
+ # smallaxe
2
+
3
+ [![CI](https://github.com/henokyemam/smallaxe/actions/workflows/ci.yml/badge.svg)](https://github.com/henokyemam/smallaxe/actions/workflows/ci.yml)
4
+
5
+ A PySpark MLOps library that simplifies model training, evaluation, and optimization for PySpark DataFrames.
6
+
7
+ ## Why smallaxe?
8
+
9
+ PySpark MLlib has a steep learning curve and verbose API. **smallaxe** provides a clean, scikit-learn-like interface for common ML workflows while leveraging the distributed power of Spark.
10
+
11
+ ## Features
12
+
13
+ - **Simple API** - Train models with familiar `fit()`/`predict()` patterns
14
+ - **Multiple Algorithms** - XGBoost, LightGBM, CatBoost, and Random Forest
15
+ - **Preprocessing Pipeline** - Imputer, Scaler, Encoder with chainable pipelines
16
+ - **Hyperparameter Optimization** - Built-in hyperopt integration with early stopping
17
+ - **Automated Training** - Train all algorithms and compare with one call
18
+ - **Visualization** - Plotly-based charts for model evaluation
19
+ - **Cross-Validation** - Train/test split and k-fold with stratified sampling
20
+
21
+ ## Installation
22
+
23
+ ```bash
24
+ pip install smallaxe
25
+ ```
26
+
27
+ Install with optional algorithm dependencies:
28
+
29
+ ```bash
30
+ pip install smallaxe[xgboost] # XGBoost support
31
+ pip install smallaxe[lightgbm] # LightGBM support
32
+ pip install smallaxe[catboost] # CatBoost support
33
+ pip install smallaxe[all] # All algorithms
34
+ ```
35
+
36
+ ## Quick Start
37
+
38
+ ```python
39
+ from smallaxe.training import Regressors
40
+ from smallaxe.datasets import load_sample_regression
41
+
42
+ # Load sample data
43
+ df = load_sample_regression(spark)
44
+
45
+ # Train a model
46
+ model = Regressors.random_forest()
47
+ model.fit(df, label_col='price', exclude_cols=['id'])
48
+
49
+ # Make predictions
50
+ predictions = model.predict(df)
51
+ ```
52
+
53
+ ## Usage Examples
54
+
55
+ ### Training with Cross-Validation
56
+
57
+ ```python
58
+ from smallaxe.training import Classifiers
59
+
60
+ model = Classifiers.xgboost(task='binary')
61
+ model.fit(
62
+ df,
63
+ label_col='churn',
64
+ validation='kfold',
65
+ n_folds=5,
66
+ stratified=True
67
+ )
68
+
69
+ print(model.validation_scores)
70
+ ```
71
+
72
+ ### Preprocessing Pipeline
73
+
74
+ ```python
75
+ from smallaxe.pipeline import Pipeline
76
+ from smallaxe.preprocessing import Imputer, Scaler, Encoder
77
+ from smallaxe.training import Regressors
78
+
79
+ pipeline = Pipeline([
80
+ ('imputer', Imputer(numerical_strategy='median')),
81
+ ('scaler', Scaler(method='standard')),
82
+ ('encoder', Encoder(method='onehot')),
83
+ ('model', Regressors.xgboost())
84
+ ])
85
+
86
+ pipeline.fit(
87
+ df,
88
+ label_col='target',
89
+ numerical_cols=['age', 'income'],
90
+ categorical_cols=['city', 'category']
91
+ )
92
+
93
+ predictions = pipeline.predict(new_df)
94
+ ```
95
+
96
+ ### Hyperparameter Optimization
97
+
98
+ ```python
99
+ from smallaxe.search import optimize
100
+ from hyperopt import hp
101
+
102
+ param_grid = {
103
+ 'max_depth': hp.choice('max_depth', [3, 5, 7, 10]),
104
+ 'learning_rate': hp.uniform('learning_rate', 0.01, 0.3)
105
+ }
106
+
107
+ best_model = optimize.run(
108
+ model=Regressors.xgboost(),
109
+ dataframe=df,
110
+ label_col='target',
111
+ param_grid=param_grid,
112
+ metric='rmse',
113
+ max_evals=50
114
+ )
115
+
116
+ print(best_model.best_params)
117
+ ```
118
+
119
+ ### Automated Training
120
+
121
+ ```python
122
+ from smallaxe.auto import AutomatedTraining
123
+
124
+ auto = AutomatedTraining(model_type='classification', metrics=['f1_score', 'auc_roc'])
125
+ auto.fit(
126
+ df,
127
+ label_col='churn',
128
+ numerical_cols=['tenure', 'monthly_charges'],
129
+ categorical_cols=['contract'],
130
+ n_folds=5
131
+ )
132
+
133
+ # Compare all models
134
+ auto.metrics.show()
135
+
136
+ # Use best model
137
+ predictions = auto.predict(new_df)
138
+ ```
139
+
140
+ ## Supported Algorithms
141
+
142
+ | Algorithm | Regressor | Classifier | Dependencies |
143
+ |-----------|-----------|------------|--------------|
144
+ | Random Forest | Yes | Yes | None (native PySpark) |
145
+ | XGBoost | Yes | Yes | `smallaxe[xgboost]` |
146
+ | LightGBM | Yes | Yes | `smallaxe[lightgbm]` |
147
+ | CatBoost | Yes | Yes | `smallaxe[catboost]` |
148
+
149
+ ## Requirements
150
+
151
+ - Python 3.8 - 3.12
152
+ - PySpark 3.3+
153
+
154
+ ## License
155
+
156
+ MIT License
@@ -0,0 +1,102 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "smallaxe"
7
+ version = "0.1.0"
8
+ description = "A PySpark MLOps library for simplified model training and optimization"
9
+ readme = "README.md"
10
+ license = {text = "MIT"}
11
+ requires-python = ">=3.8,<3.13"
12
+ authors = [
13
+ {name = "Henok Yemam"}
14
+ ]
15
+ keywords = ["pyspark", "mlops", "machine-learning", "spark", "xgboost", "lightgbm", "catboost"]
16
+ classifiers = [
17
+ "Development Status :: 3 - Alpha",
18
+ "Intended Audience :: Developers",
19
+ "Intended Audience :: Science/Research",
20
+ "License :: OSI Approved :: MIT License",
21
+ "Operating System :: OS Independent",
22
+ "Programming Language :: Python :: 3",
23
+ "Programming Language :: Python :: 3.8",
24
+ "Programming Language :: Python :: 3.9",
25
+ "Programming Language :: Python :: 3.10",
26
+ "Programming Language :: Python :: 3.11",
27
+ "Programming Language :: Python :: 3.12",
28
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
29
+ "Topic :: Software Development :: Libraries :: Python Modules",
30
+ ]
31
+
32
+ dependencies = [
33
+ "pyspark>=3.3,<4.0",
34
+ "hyperopt>=0.2.7",
35
+ "plotly>=5.0.0",
36
+ ]
37
+
38
+ [project.optional-dependencies]
39
+ xgboost = [
40
+ "xgboost>=1.7.0",
41
+ ]
42
+ lightgbm = [
43
+ "lightgbm>=3.3.0",
44
+ "synapse-ml-lightgbm>=1.0.0",
45
+ ]
46
+ catboost = [
47
+ "catboost>=1.1.0",
48
+ ]
49
+ all = [
50
+ "smallaxe[xgboost]",
51
+ "smallaxe[lightgbm]",
52
+ "smallaxe[catboost]",
53
+ ]
54
+ dev = [
55
+ "pytest>=7.0.0",
56
+ "pytest-spark>=0.6.0",
57
+ "black>=23.0.0",
58
+ "ruff>=0.1.0",
59
+ "mypy>=1.0.0",
60
+ ]
61
+
62
+ [project.urls]
63
+ Homepage = "https://github.com/henokyemam/smallaxe"
64
+ Repository = "https://github.com/henokyemam/smallaxe"
65
+ Issues = "https://github.com/henokyemam/smallaxe/issues"
66
+
67
+ [tool.setuptools.packages.find]
68
+ where = ["."]
69
+ include = ["smallaxe*"]
70
+
71
+ [tool.black]
72
+ line-length = 100
73
+ target-version = ["py38", "py39", "py310", "py311", "py312"]
74
+
75
+ [tool.ruff]
76
+ line-length = 100
77
+ target-version = "py38"
78
+
79
+ [tool.ruff.lint]
80
+ select = [
81
+ "E", # pycodestyle errors
82
+ "W", # pycodestyle warnings
83
+ "F", # pyflakes
84
+ "I", # isort
85
+ "B", # flake8-bugbear
86
+ "UP", # pyupgrade
87
+ ]
88
+ ignore = [
89
+ "E501", # line too long (handled by black)
90
+ ]
91
+
92
+ [tool.mypy]
93
+ python_version = "3.8"
94
+ warn_return_any = true
95
+ warn_unused_configs = true
96
+ ignore_missing_imports = true
97
+
98
+ [tool.pytest.ini_options]
99
+ testpaths = ["tests"]
100
+ python_files = ["test_*.py"]
101
+ python_functions = ["test_*"]
102
+ addopts = "-v"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,157 @@
1
+ """
2
+ smallaxe - A PySpark MLOps library for simplified model training and optimization.
3
+ """
4
+
5
+ from contextlib import contextmanager
6
+ from typing import Any, Generator, Optional
7
+
8
+ from smallaxe._config import (
9
+ VALID_CACHE_STRATEGIES,
10
+ VALID_VERBOSITY_LEVELS,
11
+ _config,
12
+ )
13
+ from smallaxe.exceptions import ConfigurationError
14
+
15
+ __version__ = "0.1.0"
16
+
17
+ __all__ = [
18
+ "__version__",
19
+ "set_verbosity",
20
+ "get_verbosity",
21
+ "verbosity",
22
+ "set_spark_session",
23
+ "get_spark_session",
24
+ "set_seed",
25
+ "get_seed",
26
+ "set_cache_strategy",
27
+ "get_cache_strategy",
28
+ ]
29
+
30
+
31
+ def set_verbosity(level: str) -> None:
32
+ """Set the global verbosity level.
33
+
34
+ Args:
35
+ level: Verbosity level. One of 'quiet', 'normal', or 'verbose'.
36
+ - 'quiet': Only errors, no progress bars or info messages
37
+ - 'normal': Progress bars and key info (default)
38
+ - 'verbose': Detailed logging for debugging
39
+
40
+ Raises:
41
+ ConfigurationError: If level is not a valid verbosity level.
42
+ """
43
+ if level not in VALID_VERBOSITY_LEVELS:
44
+ raise ConfigurationError(
45
+ setting="verbosity",
46
+ value=level,
47
+ allowed_values=list(VALID_VERBOSITY_LEVELS),
48
+ )
49
+ _config._verbosity = level
50
+
51
+
52
+ def get_verbosity() -> str:
53
+ """Get the current verbosity level.
54
+
55
+ Returns:
56
+ The current verbosity level ('quiet', 'normal', or 'verbose').
57
+ """
58
+ return _config._verbosity
59
+
60
+
61
+ @contextmanager
62
+ def verbosity(level: str) -> Generator[None, None, None]:
63
+ """Context manager for temporarily changing verbosity level.
64
+
65
+ Args:
66
+ level: Verbosity level to use within the context.
67
+
68
+ Raises:
69
+ ConfigurationError: If level is not a valid verbosity level.
70
+
71
+ Example:
72
+ >>> with smallaxe.verbosity('quiet'):
73
+ ... model.fit(df, label_col='target') # runs silently
74
+ """
75
+ previous_level = get_verbosity()
76
+ set_verbosity(level)
77
+ try:
78
+ yield
79
+ finally:
80
+ _config._verbosity = previous_level
81
+
82
+
83
+ def set_spark_session(spark: Any) -> None:
84
+ """Set the Spark session to use.
85
+
86
+ Args:
87
+ spark: A SparkSession instance. If None, smallaxe will attempt
88
+ to get or create a session when needed.
89
+ """
90
+ _config._spark_session = spark
91
+
92
+
93
+ def get_spark_session() -> Optional[Any]:
94
+ """Get the configured Spark session.
95
+
96
+ Returns:
97
+ The configured SparkSession, or None if not set.
98
+ """
99
+ return _config._spark_session
100
+
101
+
102
+ def set_seed(seed: Optional[int]) -> None:
103
+ """Set the global random seed for reproducibility.
104
+
105
+ This affects all random operations including train/test splits,
106
+ k-fold cross-validation, and hyperopt sampling.
107
+
108
+ Args:
109
+ seed: Integer seed value, or None to reset to no seed.
110
+
111
+ Raises:
112
+ ConfigurationError: If seed is not an integer or None.
113
+ """
114
+ if seed is not None and not isinstance(seed, int):
115
+ raise ConfigurationError(
116
+ message=f"Seed must be an integer or None, got {type(seed).__name__}."
117
+ )
118
+ _config._seed = seed
119
+
120
+
121
+ def get_seed() -> Optional[int]:
122
+ """Get the current random seed.
123
+
124
+ Returns:
125
+ The current seed value, or None if not set.
126
+ """
127
+ return _config._seed
128
+
129
+
130
+ def set_cache_strategy(strategy: str) -> None:
131
+ """Set the caching strategy for PySpark operations.
132
+
133
+ Args:
134
+ strategy: Cache strategy. One of 'auto', 'always', or 'never'.
135
+ - 'auto': Smart caching - cache after preprocessing, unpersist after training
136
+ - 'always': Cache at every stage (use for debugging or small datasets)
137
+ - 'never': No automatic caching (manual control)
138
+
139
+ Raises:
140
+ ConfigurationError: If strategy is not a valid cache strategy.
141
+ """
142
+ if strategy not in VALID_CACHE_STRATEGIES:
143
+ raise ConfigurationError(
144
+ setting="cache_strategy",
145
+ value=strategy,
146
+ allowed_values=list(VALID_CACHE_STRATEGIES),
147
+ )
148
+ _config._cache_strategy = strategy
149
+
150
+
151
+ def get_cache_strategy() -> str:
152
+ """Get the current cache strategy.
153
+
154
+ Returns:
155
+ The current cache strategy ('auto', 'always', or 'never').
156
+ """
157
+ return _config._cache_strategy
@@ -0,0 +1,37 @@
1
+ """Internal configuration state for smallaxe."""
2
+
3
+ from typing import Any, Optional
4
+
5
+ # Valid configuration values
6
+ VALID_VERBOSITY_LEVELS = ("quiet", "normal", "verbose")
7
+ VALID_CACHE_STRATEGIES = ("auto", "always", "never")
8
+
9
+ # Default configuration values
10
+ DEFAULT_VERBOSITY = "normal"
11
+ DEFAULT_CACHE_STRATEGY = "auto"
12
+ DEFAULT_SEED = None
13
+
14
+
15
+ class _Config:
16
+ """Internal configuration state container.
17
+
18
+ This class holds the global configuration state for smallaxe.
19
+ It should not be accessed directly - use the module-level functions instead.
20
+ """
21
+
22
+ def __init__(self) -> None:
23
+ self._verbosity: str = DEFAULT_VERBOSITY
24
+ self._cache_strategy: str = DEFAULT_CACHE_STRATEGY
25
+ self._seed: Optional[int] = DEFAULT_SEED
26
+ self._spark_session: Optional[Any] = None
27
+
28
+ def reset(self) -> None:
29
+ """Reset configuration to defaults."""
30
+ self._verbosity = DEFAULT_VERBOSITY
31
+ self._cache_strategy = DEFAULT_CACHE_STRATEGY
32
+ self._seed = DEFAULT_SEED
33
+ self._spark_session = None
34
+
35
+
36
+ # Global configuration instance
37
+ _config = _Config()
@@ -0,0 +1 @@
1
+ """Auto module - automated training."""
@@ -0,0 +1,13 @@
1
+ """Sample datasets for testing and demos."""
2
+
3
+ from smallaxe.datasets._data import (
4
+ dataset_info,
5
+ load_sample_classification,
6
+ load_sample_regression,
7
+ )
8
+
9
+ __all__ = [
10
+ "load_sample_regression",
11
+ "load_sample_classification",
12
+ "dataset_info",
13
+ ]