mostlyai-engine 1.7.1__tar.gz → 2.0.1__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 (55) hide show
  1. mostlyai_engine-2.0.1/PKG-INFO +323 -0
  2. mostlyai_engine-2.0.1/README.md +272 -0
  3. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/__init__.py +14 -2
  4. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_common.py +115 -0
  5. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/numeric.py +22 -15
  6. mostlyai_engine-2.0.1/mostlyai/engine/_language/interface.py +322 -0
  7. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/encoding.py +5 -1
  8. mostlyai_engine-2.0.1/mostlyai/engine/_tabular/interface.py +657 -0
  9. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/logging.py +14 -0
  10. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/pyproject.toml +4 -2
  11. mostlyai_engine-1.7.1/PKG-INFO +0 -205
  12. mostlyai_engine-1.7.1/README.md +0 -155
  13. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/.gitignore +0 -0
  14. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/LICENSE +0 -0
  15. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_dtypes.py +0 -0
  16. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/__init__.py +0 -0
  17. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/language/__init__.py +0 -0
  18. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/language/categorical.py +0 -0
  19. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/language/datetime.py +0 -0
  20. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/language/numeric.py +0 -0
  21. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/language/text.py +0 -0
  22. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/__init__.py +0 -0
  23. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/categorical.py +0 -0
  24. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/character.py +0 -0
  25. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/datetime.py +0 -0
  26. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/itt.py +0 -0
  27. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_encoding_types/tabular/lat_long.py +0 -0
  28. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/__init__.py +0 -0
  29. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/common.py +0 -0
  30. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/encoding.py +0 -0
  31. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/engine/__init__.py +0 -0
  32. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/engine/base.py +0 -0
  33. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/engine/hf_engine.py +0 -0
  34. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/engine/vllm_engine.py +0 -0
  35. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/generation.py +0 -0
  36. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/lstm.py +0 -0
  37. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/tokenizer_utils.py +0 -0
  38. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/training.py +0 -0
  39. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_language/xgrammar_utils.py +0 -0
  40. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_memory.py +0 -0
  41. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/__init__.py +0 -0
  42. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/argn.py +0 -0
  43. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/common.py +0 -0
  44. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/fairness.py +0 -0
  45. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/generation.py +0 -0
  46. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_tabular/training.py +0 -0
  47. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_training_utils.py +0 -0
  48. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/_workspace.py +0 -0
  49. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/analysis.py +0 -0
  50. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/domain.py +0 -0
  51. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/encoding.py +0 -0
  52. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/generation.py +0 -0
  53. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/random_state.py +0 -0
  54. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/splitting.py +0 -0
  55. {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.1}/mostlyai/engine/training.py +0 -0
@@ -0,0 +1,323 @@
1
+ Metadata-Version: 2.4
2
+ Name: mostlyai-engine
3
+ Version: 2.0.1
4
+ Summary: Synthetic Data Engine
5
+ Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
6
+ Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine
7
+ Project-URL: documentation, https://mostly-ai.github.io/mostlyai-engine/
8
+ Author-email: MOSTLY AI <dev@mostly.ai>
9
+ License-Expression: Apache-2.0
10
+ License-File: LICENSE
11
+ Classifier: Development Status :: 5 - Production/Stable
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Financial and Insurance Industry
14
+ Classifier: Intended Audience :: Healthcare Industry
15
+ Classifier: Intended Audience :: Information Technology
16
+ Classifier: Intended Audience :: Science/Research
17
+ Classifier: Intended Audience :: Telecommunications Industry
18
+ Classifier: License :: OSI Approved :: Apache Software License
19
+ Classifier: Operating System :: OS Independent
20
+ Classifier: Programming Language :: Python :: 3.10
21
+ Classifier: Programming Language :: Python :: 3.11
22
+ Classifier: Programming Language :: Python :: 3.12
23
+ Classifier: Programming Language :: Python :: 3.13
24
+ Classifier: Topic :: Software Development :: Libraries
25
+ Classifier: Typing :: Typed
26
+ Requires-Python: >=3.10
27
+ Requires-Dist: accelerate>=1.5.0
28
+ Requires-Dist: datasets>=3.0.0
29
+ Requires-Dist: huggingface-hub[hf-xet]>=0.30.2
30
+ Requires-Dist: joblib>=1.4.2
31
+ Requires-Dist: json-repair>=0.47.0
32
+ Requires-Dist: numpy>=2.0.0
33
+ Requires-Dist: opacus>=1.5.4
34
+ Requires-Dist: pandas>=2.2.0
35
+ Requires-Dist: peft>=0.12.0
36
+ Requires-Dist: psutil<6,>=5.9.5
37
+ Requires-Dist: pyarrow>=16.0.0
38
+ Requires-Dist: scikit-learn>=1.4.0
39
+ Requires-Dist: setuptools>=77.0.3
40
+ Requires-Dist: tokenizers>=0.21.0
41
+ Requires-Dist: torch<2.9.0,>=2.8.0
42
+ Requires-Dist: torchaudio<2.9.0,>=2.8.0
43
+ Requires-Dist: torchvision<0.24.0,>=0.23.0
44
+ Requires-Dist: transformers>=4.55.0
45
+ Requires-Dist: xgrammar>=0.1.21
46
+ Provides-Extra: gpu
47
+ Requires-Dist: bitsandbytes==0.42.0; (sys_platform == 'darwin') and extra == 'gpu'
48
+ Requires-Dist: bitsandbytes>=0.45.5; (sys_platform == 'linux') and extra == 'gpu'
49
+ Requires-Dist: vllm==0.10.2; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
50
+ Description-Content-Type: text/markdown
51
+
52
+ # Synthetic Data Engine 💎
53
+
54
+ ![GitHub Release](https://img.shields.io/github/v/release/mostly-ai/mostlyai-engine)
55
+ [![Documentation](https://img.shields.io/badge/docs-latest-green)](https://mostly-ai.github.io/mostlyai-engine/)
56
+ [![stats](https://pepy.tech/badge/mostlyai-engine)](https://pypi.org/project/mostlyai-engine/)
57
+ ![license](https://img.shields.io/github/license/mostly-ai/mostlyai-engine)
58
+ ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mostlyai-engine)
59
+
60
+ [Documentation](https://mostly-ai.github.io/mostlyai-engine/) | [Technical Paper](https://arxiv.org/abs/2501.12012) | [Free Cloud Service](https://app.mostly.ai/)
61
+
62
+ Create high-fidelity privacy-safe synthetic data:
63
+
64
+ 1. train a generative model once:
65
+ * train on flat or sequential data
66
+ * control training time & params
67
+ * monitor training progress
68
+ * optionally enable differential privacy
69
+ * optionally provide context data
70
+ 2. generate synthetic data samples to your needs:
71
+ * up-sample / down-sample
72
+ * conditionally generate
73
+ * rebalance categories
74
+ * impute missing values
75
+ * incorporate fairness
76
+ * adjust sampling temperature
77
+ * predict / classify / regress
78
+ * detect outliers / anomalies
79
+ * and more
80
+
81
+ ...all within your own compute environment, all with a few lines of Python code 💥.
82
+
83
+ Note: Models only need to be trained once and can then be flexibly reused for various downstream tasks — such as regression, classification, imputation, or sampling — without the need for retraining.
84
+
85
+ Two model classes with these methods are available:
86
+
87
+ 1. `TabularARGN()`: For structured, flat or sequential tabular data.
88
+ * `argn.fit(data)`: Train a TabularARGN model
89
+ * `argn.sample(n_samples)`: Generate samples
90
+ * `argn.predict(target, n_draws, agg_fn)`: Predict a feature
91
+ * `argn.predict_proba(target, n_draws)`: Estimate probabilities
92
+ * `argn.impute(data)`: Fill missing values
93
+ 2. `LanguageModel()`: For semi-structured, flat textual tabular data.
94
+ * `.fit(data)`: Train a Language model
95
+ * `.sample(n_samples)`: Generate samples
96
+
97
+ This library serves as the core model engine for the [Synthetic Data SDK](https://github.com/mostly-ai/mostlyai). For an easy-to-use, higher-level toolkit, please refer to the SDK.
98
+
99
+
100
+ ## Installation
101
+
102
+ It is highly recommended to install the package within a dedicated virtual environment using [uv](https://docs.astral.sh/uv/).
103
+
104
+ The latest release of `mostlyai-engine` can be installed via uv:
105
+
106
+ ```bash
107
+ uv pip install -U mostlyai-engine
108
+ ```
109
+
110
+ or alternatively for a GPU setup (needed for LLM finetuning and inference):
111
+ ```bash
112
+ uv pip install -U 'mostlyai-engine[gpu]'
113
+ ```
114
+
115
+ On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
116
+
117
+ ```bash
118
+ uv pip install -U torch==2.8.0+cpu torchvision==0.23.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
119
+ ```
120
+
121
+ ## TabularARGN for Flat Data
122
+
123
+ The `TabularARGN` class provides a scikit-learn-compatible interface for working with structured tabular data. It can be used for synthetic data generation, classification, regression, and imputation.
124
+
125
+ ### Model Training
126
+
127
+ Load your data and train the model:
128
+
129
+ ```python
130
+ import pandas as pd
131
+ from sklearn.model_selection import train_test_split
132
+ from mostlyai.engine import TabularARGN
133
+
134
+ # prepare data
135
+ data = pd.read_csv("https://github.com/user-attachments/files/23480587/census10k.csv.gz")
136
+ data_train, data_test = train_test_split(data, test_size=0.2)
137
+
138
+ # fit TabularARGN
139
+ argn = TabularARGN()
140
+ argn.fit(data_train)
141
+ ```
142
+
143
+ ### Sampling / Synthetic Data Generation
144
+
145
+ Generate new synthetic samples:
146
+
147
+ ```python
148
+ # unconditional sampling
149
+ argn.sample(n_samples=1000)
150
+ ```
151
+
152
+ Generate new synthetic samples conditionally:
153
+
154
+ ```python
155
+ # prepare seed
156
+ seed_data = pd.DataFrame({
157
+ "age": [25, 50],
158
+ "education": ["Bachelors", "HS-grad"]
159
+ })
160
+
161
+ # conditional sampling
162
+ argn.sample(seed_data=seed_data)
163
+ ```
164
+
165
+ ### Imputation / Filling Gaps
166
+
167
+ Fill in missing values:
168
+
169
+ ```python
170
+ # prepare demo data with missings
171
+ data_with_missings = data_test.head(300).reset_index(drop=True)
172
+ data_with_missings.loc[0:299, "age"] = pd.NA
173
+ data_with_missings.loc[0:199, "race"] = pd.NA
174
+ data_with_missings.loc[100:299, "income"] = pd.NA
175
+
176
+ # impute missing values each with a random sample
177
+ data_imputed = argn.impute(data_with_missings)
178
+
179
+ # impute missing values each with their point estimates
180
+ data_imputed = argn.impute(data_with_missings, n_draws=100)
181
+
182
+ ```
183
+
184
+ ### Predictions / Classification
185
+
186
+ Predict any categorical target column:
187
+
188
+ ```python
189
+ from sklearn.metrics import accuracy_score, roc_auc_score
190
+
191
+ # predict class labels for a categorical
192
+ predictions = argn.predict(data_test, target="income", n_draws=100, agg_fn="mode")
193
+
194
+ # predict class probabilities for a categorical
195
+ probabilities = argn.predict_proba(data_test, target="income", n_draws=100)
196
+
197
+ # evaluate performance
198
+ accuracy = accuracy_score(data_test["income"], predictions)
199
+ auc = roc_auc_score(data_test["income"], probabilities[:, 1])
200
+ print(f"Accuracy: {accuracy:.3f}, AUC: {auc:.3f}")
201
+ ```
202
+
203
+ ### Predictions / Regression
204
+
205
+ Predict any numerical target column:
206
+
207
+ ```python
208
+ from sklearn.metrics import mean_absolute_error
209
+
210
+ # predict target values
211
+ predictions = argn.predict(data_test, target="age", n_draws=10, agg_fn="mean")
212
+
213
+ # evaluate performance
214
+ mae = mean_absolute_error(data_test["age"], predictions)
215
+ print(f"MAE: {mae:.1f} years")
216
+ ```
217
+
218
+ ## TabularARGN for Sequential Data
219
+
220
+ For sequential data (e.g., time series or event logs), specify the context key:
221
+
222
+ ### Model Training - With Context Data
223
+
224
+ ```python
225
+ import pandas as pd
226
+ from mostlyai.engine import TabularARGN
227
+
228
+ # load sequential data
229
+ tgt_data = pd.read_csv("https://github.com/user-attachments/files/23480787/batting.csv.gz")
230
+ ctx_data = pd.read_csv("https://github.com/user-attachments/files/23480786/players.csv.gz")
231
+
232
+ # fit TabularARGN with a context key column
233
+ argn = TabularARGN(
234
+ tgt_context_key="players_id",
235
+ ctx_primary_key="id",
236
+ ctx_data=ctx_data,
237
+ max_training_time=2, # 2 minutes
238
+ verbose=0,
239
+ )
240
+ argn.fit(tgt_data)
241
+ ```
242
+
243
+ ### Sampling / Synthetic Data Generation
244
+
245
+ Generate new synthetic samples (using existing context):
246
+ ```python
247
+ argn.sample(n_samples=5)
248
+ ```
249
+
250
+ Generate new synthetic samples conditionally (using custom context and seed):
251
+
252
+ ```python
253
+ ctx_data = pd.DataFrame({
254
+ "id": ["Player1", "Player2"],
255
+ "weight": [170, 160],
256
+ "height": [70, 68],
257
+ "bats": ["R", "L"],
258
+ "throws": ["R", "L"],
259
+ })
260
+ argn.sample(ctx_data=ctx_data)
261
+ ```
262
+
263
+ ## Basic Usage of LanguageModel
264
+
265
+ The `LanguageModel` class provides a scikit-learn-compatible interface for working with semi-structured textual data. It leverages pre-trained language models or trains lightweight LSTM models from scratch to generate synthetic text data.
266
+
267
+ **Note**: The default model is `MOSTLY_AI/LSTMFromScratch-3m`, a lightweight LSTM model trained from scratch (GPU strongly recommended). You can also use pre-trained HuggingFace models by setting model to e.g. `microsoft/phi-1.5` (GPU required).
268
+
269
+ ### Model Training
270
+
271
+ Load your data and train the model:
272
+
273
+ ```python
274
+ import pandas as pd
275
+ from mostlyai.engine import LanguageModel
276
+
277
+ # load data
278
+ data = pd.read_csv("https://github.com/user-attachments/files/23486562/airbnb20k.csv.gz")
279
+
280
+ # fit LanguageModel
281
+ lm = LanguageModel(
282
+ model="MOSTLY_AI/LSTMFromScratch-3m",
283
+ tgt_encoding_types={
284
+ 'neighbourhood': 'LANGUAGE_CATEGORICAL',
285
+ 'title': 'LANGUAGE_TEXT',
286
+ },
287
+ max_training_time=10, # 10 minutes
288
+ verbose=1,
289
+ )
290
+ lm.fit(data)
291
+ ```
292
+
293
+ ### Sampling / Synthetic Text Generation
294
+
295
+ Generate new synthetic samples using the trained language model:
296
+
297
+ ```python
298
+ # unconditional sampling
299
+ lm.sample(
300
+ n_samples=100,
301
+ sampling_temperature=0.8,
302
+ )
303
+ ```
304
+
305
+ ```python
306
+ # prepare seed
307
+ seed_data = pd.DataFrame({
308
+ "neighbourhood": ["Westminster", "Hackney"],
309
+ })
310
+
311
+ # conditional sampling with seed values
312
+ lm.sample(
313
+ seed_data=seed_data,
314
+ sampling_temperature=0.8,
315
+ )
316
+ ```
317
+
318
+ ## Further Examples
319
+
320
+ Example notebooks demonstrating various use cases are available in the `examples` directory:
321
+ - TabularARGN for flat tabular data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/flat.ipynb)
322
+ - TabularARGN for sequential data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)
323
+ - LanguageModel for textual data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/sequential.ipynb)
@@ -0,0 +1,272 @@
1
+ # Synthetic Data Engine 💎
2
+
3
+ ![GitHub Release](https://img.shields.io/github/v/release/mostly-ai/mostlyai-engine)
4
+ [![Documentation](https://img.shields.io/badge/docs-latest-green)](https://mostly-ai.github.io/mostlyai-engine/)
5
+ [![stats](https://pepy.tech/badge/mostlyai-engine)](https://pypi.org/project/mostlyai-engine/)
6
+ ![license](https://img.shields.io/github/license/mostly-ai/mostlyai-engine)
7
+ ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mostlyai-engine)
8
+
9
+ [Documentation](https://mostly-ai.github.io/mostlyai-engine/) | [Technical Paper](https://arxiv.org/abs/2501.12012) | [Free Cloud Service](https://app.mostly.ai/)
10
+
11
+ Create high-fidelity privacy-safe synthetic data:
12
+
13
+ 1. train a generative model once:
14
+ * train on flat or sequential data
15
+ * control training time & params
16
+ * monitor training progress
17
+ * optionally enable differential privacy
18
+ * optionally provide context data
19
+ 2. generate synthetic data samples to your needs:
20
+ * up-sample / down-sample
21
+ * conditionally generate
22
+ * rebalance categories
23
+ * impute missing values
24
+ * incorporate fairness
25
+ * adjust sampling temperature
26
+ * predict / classify / regress
27
+ * detect outliers / anomalies
28
+ * and more
29
+
30
+ ...all within your own compute environment, all with a few lines of Python code 💥.
31
+
32
+ Note: Models only need to be trained once and can then be flexibly reused for various downstream tasks — such as regression, classification, imputation, or sampling — without the need for retraining.
33
+
34
+ Two model classes with these methods are available:
35
+
36
+ 1. `TabularARGN()`: For structured, flat or sequential tabular data.
37
+ * `argn.fit(data)`: Train a TabularARGN model
38
+ * `argn.sample(n_samples)`: Generate samples
39
+ * `argn.predict(target, n_draws, agg_fn)`: Predict a feature
40
+ * `argn.predict_proba(target, n_draws)`: Estimate probabilities
41
+ * `argn.impute(data)`: Fill missing values
42
+ 2. `LanguageModel()`: For semi-structured, flat textual tabular data.
43
+ * `.fit(data)`: Train a Language model
44
+ * `.sample(n_samples)`: Generate samples
45
+
46
+ This library serves as the core model engine for the [Synthetic Data SDK](https://github.com/mostly-ai/mostlyai). For an easy-to-use, higher-level toolkit, please refer to the SDK.
47
+
48
+
49
+ ## Installation
50
+
51
+ It is highly recommended to install the package within a dedicated virtual environment using [uv](https://docs.astral.sh/uv/).
52
+
53
+ The latest release of `mostlyai-engine` can be installed via uv:
54
+
55
+ ```bash
56
+ uv pip install -U mostlyai-engine
57
+ ```
58
+
59
+ or alternatively for a GPU setup (needed for LLM finetuning and inference):
60
+ ```bash
61
+ uv pip install -U 'mostlyai-engine[gpu]'
62
+ ```
63
+
64
+ On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
65
+
66
+ ```bash
67
+ uv pip install -U torch==2.8.0+cpu torchvision==0.23.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
68
+ ```
69
+
70
+ ## TabularARGN for Flat Data
71
+
72
+ The `TabularARGN` class provides a scikit-learn-compatible interface for working with structured tabular data. It can be used for synthetic data generation, classification, regression, and imputation.
73
+
74
+ ### Model Training
75
+
76
+ Load your data and train the model:
77
+
78
+ ```python
79
+ import pandas as pd
80
+ from sklearn.model_selection import train_test_split
81
+ from mostlyai.engine import TabularARGN
82
+
83
+ # prepare data
84
+ data = pd.read_csv("https://github.com/user-attachments/files/23480587/census10k.csv.gz")
85
+ data_train, data_test = train_test_split(data, test_size=0.2)
86
+
87
+ # fit TabularARGN
88
+ argn = TabularARGN()
89
+ argn.fit(data_train)
90
+ ```
91
+
92
+ ### Sampling / Synthetic Data Generation
93
+
94
+ Generate new synthetic samples:
95
+
96
+ ```python
97
+ # unconditional sampling
98
+ argn.sample(n_samples=1000)
99
+ ```
100
+
101
+ Generate new synthetic samples conditionally:
102
+
103
+ ```python
104
+ # prepare seed
105
+ seed_data = pd.DataFrame({
106
+ "age": [25, 50],
107
+ "education": ["Bachelors", "HS-grad"]
108
+ })
109
+
110
+ # conditional sampling
111
+ argn.sample(seed_data=seed_data)
112
+ ```
113
+
114
+ ### Imputation / Filling Gaps
115
+
116
+ Fill in missing values:
117
+
118
+ ```python
119
+ # prepare demo data with missings
120
+ data_with_missings = data_test.head(300).reset_index(drop=True)
121
+ data_with_missings.loc[0:299, "age"] = pd.NA
122
+ data_with_missings.loc[0:199, "race"] = pd.NA
123
+ data_with_missings.loc[100:299, "income"] = pd.NA
124
+
125
+ # impute missing values each with a random sample
126
+ data_imputed = argn.impute(data_with_missings)
127
+
128
+ # impute missing values each with their point estimates
129
+ data_imputed = argn.impute(data_with_missings, n_draws=100)
130
+
131
+ ```
132
+
133
+ ### Predictions / Classification
134
+
135
+ Predict any categorical target column:
136
+
137
+ ```python
138
+ from sklearn.metrics import accuracy_score, roc_auc_score
139
+
140
+ # predict class labels for a categorical
141
+ predictions = argn.predict(data_test, target="income", n_draws=100, agg_fn="mode")
142
+
143
+ # predict class probabilities for a categorical
144
+ probabilities = argn.predict_proba(data_test, target="income", n_draws=100)
145
+
146
+ # evaluate performance
147
+ accuracy = accuracy_score(data_test["income"], predictions)
148
+ auc = roc_auc_score(data_test["income"], probabilities[:, 1])
149
+ print(f"Accuracy: {accuracy:.3f}, AUC: {auc:.3f}")
150
+ ```
151
+
152
+ ### Predictions / Regression
153
+
154
+ Predict any numerical target column:
155
+
156
+ ```python
157
+ from sklearn.metrics import mean_absolute_error
158
+
159
+ # predict target values
160
+ predictions = argn.predict(data_test, target="age", n_draws=10, agg_fn="mean")
161
+
162
+ # evaluate performance
163
+ mae = mean_absolute_error(data_test["age"], predictions)
164
+ print(f"MAE: {mae:.1f} years")
165
+ ```
166
+
167
+ ## TabularARGN for Sequential Data
168
+
169
+ For sequential data (e.g., time series or event logs), specify the context key:
170
+
171
+ ### Model Training - With Context Data
172
+
173
+ ```python
174
+ import pandas as pd
175
+ from mostlyai.engine import TabularARGN
176
+
177
+ # load sequential data
178
+ tgt_data = pd.read_csv("https://github.com/user-attachments/files/23480787/batting.csv.gz")
179
+ ctx_data = pd.read_csv("https://github.com/user-attachments/files/23480786/players.csv.gz")
180
+
181
+ # fit TabularARGN with a context key column
182
+ argn = TabularARGN(
183
+ tgt_context_key="players_id",
184
+ ctx_primary_key="id",
185
+ ctx_data=ctx_data,
186
+ max_training_time=2, # 2 minutes
187
+ verbose=0,
188
+ )
189
+ argn.fit(tgt_data)
190
+ ```
191
+
192
+ ### Sampling / Synthetic Data Generation
193
+
194
+ Generate new synthetic samples (using existing context):
195
+ ```python
196
+ argn.sample(n_samples=5)
197
+ ```
198
+
199
+ Generate new synthetic samples conditionally (using custom context and seed):
200
+
201
+ ```python
202
+ ctx_data = pd.DataFrame({
203
+ "id": ["Player1", "Player2"],
204
+ "weight": [170, 160],
205
+ "height": [70, 68],
206
+ "bats": ["R", "L"],
207
+ "throws": ["R", "L"],
208
+ })
209
+ argn.sample(ctx_data=ctx_data)
210
+ ```
211
+
212
+ ## Basic Usage of LanguageModel
213
+
214
+ The `LanguageModel` class provides a scikit-learn-compatible interface for working with semi-structured textual data. It leverages pre-trained language models or trains lightweight LSTM models from scratch to generate synthetic text data.
215
+
216
+ **Note**: The default model is `MOSTLY_AI/LSTMFromScratch-3m`, a lightweight LSTM model trained from scratch (GPU strongly recommended). You can also use pre-trained HuggingFace models by setting model to e.g. `microsoft/phi-1.5` (GPU required).
217
+
218
+ ### Model Training
219
+
220
+ Load your data and train the model:
221
+
222
+ ```python
223
+ import pandas as pd
224
+ from mostlyai.engine import LanguageModel
225
+
226
+ # load data
227
+ data = pd.read_csv("https://github.com/user-attachments/files/23486562/airbnb20k.csv.gz")
228
+
229
+ # fit LanguageModel
230
+ lm = LanguageModel(
231
+ model="MOSTLY_AI/LSTMFromScratch-3m",
232
+ tgt_encoding_types={
233
+ 'neighbourhood': 'LANGUAGE_CATEGORICAL',
234
+ 'title': 'LANGUAGE_TEXT',
235
+ },
236
+ max_training_time=10, # 10 minutes
237
+ verbose=1,
238
+ )
239
+ lm.fit(data)
240
+ ```
241
+
242
+ ### Sampling / Synthetic Text Generation
243
+
244
+ Generate new synthetic samples using the trained language model:
245
+
246
+ ```python
247
+ # unconditional sampling
248
+ lm.sample(
249
+ n_samples=100,
250
+ sampling_temperature=0.8,
251
+ )
252
+ ```
253
+
254
+ ```python
255
+ # prepare seed
256
+ seed_data = pd.DataFrame({
257
+ "neighbourhood": ["Westminster", "Hackney"],
258
+ })
259
+
260
+ # conditional sampling with seed values
261
+ lm.sample(
262
+ seed_data=seed_data,
263
+ sampling_temperature=0.8,
264
+ )
265
+ ```
266
+
267
+ ## Further Examples
268
+
269
+ Example notebooks demonstrating various use cases are available in the `examples` directory:
270
+ - TabularARGN for flat tabular data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/flat.ipynb)
271
+ - TabularARGN for sequential data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)
272
+ - LanguageModel for textual data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/sequential.ipynb)
@@ -13,6 +13,8 @@
13
13
  # limitations under the License.
14
14
  import warnings
15
15
 
16
+ from mostlyai.engine._language.interface import LanguageModel
17
+ from mostlyai.engine._tabular.interface import TabularARGN
16
18
  from mostlyai.engine.analysis import analyze
17
19
  from mostlyai.engine.encoding import encode
18
20
  from mostlyai.engine.generation import generate
@@ -21,8 +23,18 @@ from mostlyai.engine.random_state import set_random_state
21
23
  from mostlyai.engine.splitting import split
22
24
  from mostlyai.engine.training import train
23
25
 
24
- __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging", "set_random_state"]
25
- __version__ = "1.7.1"
26
+ __all__ = [
27
+ "split",
28
+ "analyze",
29
+ "encode",
30
+ "train",
31
+ "generate",
32
+ "init_logging",
33
+ "set_random_state",
34
+ "TabularARGN",
35
+ "LanguageModel",
36
+ ]
37
+ __version__ = "2.0.1"
26
38
 
27
39
  # suppress specific warning related to os.fork() in multi-threaded processes
28
40
  warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")