mostlyai-engine 1.7.0__tar.gz → 2.0.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 (55) hide show
  1. mostlyai_engine-2.0.0/PKG-INFO +319 -0
  2. mostlyai_engine-2.0.0/README.md +268 -0
  3. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/__init__.py +14 -2
  4. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_common.py +115 -0
  5. mostlyai_engine-2.0.0/mostlyai/engine/_language/interface.py +322 -0
  6. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/training.py +1 -7
  7. mostlyai_engine-2.0.0/mostlyai/engine/_tabular/interface.py +657 -0
  8. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/training.py +1 -7
  9. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/logging.py +14 -0
  10. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/pyproject.toml +4 -2
  11. mostlyai_engine-1.7.0/PKG-INFO +0 -205
  12. mostlyai_engine-1.7.0/README.md +0 -155
  13. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/.gitignore +0 -0
  14. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/LICENSE +0 -0
  15. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_dtypes.py +0 -0
  16. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/__init__.py +0 -0
  17. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/__init__.py +0 -0
  18. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/categorical.py +0 -0
  19. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/datetime.py +0 -0
  20. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/numeric.py +0 -0
  21. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/text.py +0 -0
  22. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/__init__.py +0 -0
  23. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/categorical.py +0 -0
  24. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/character.py +0 -0
  25. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/datetime.py +0 -0
  26. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/itt.py +0 -0
  27. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/lat_long.py +0 -0
  28. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/numeric.py +0 -0
  29. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/__init__.py +0 -0
  30. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/common.py +0 -0
  31. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/encoding.py +0 -0
  32. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/__init__.py +0 -0
  33. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/base.py +0 -0
  34. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/hf_engine.py +0 -0
  35. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/vllm_engine.py +0 -0
  36. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/generation.py +0 -0
  37. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/lstm.py +0 -0
  38. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/tokenizer_utils.py +0 -0
  39. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/xgrammar_utils.py +0 -0
  40. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_memory.py +0 -0
  41. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/__init__.py +0 -0
  42. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/argn.py +0 -0
  43. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/common.py +0 -0
  44. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/encoding.py +0 -0
  45. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/fairness.py +0 -0
  46. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/generation.py +0 -0
  47. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_training_utils.py +0 -0
  48. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/_workspace.py +0 -0
  49. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/analysis.py +0 -0
  50. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/domain.py +0 -0
  51. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/encoding.py +0 -0
  52. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/generation.py +0 -0
  53. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/random_state.py +0 -0
  54. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/splitting.py +0 -0
  55. {mostlyai_engine-1.7.0 → mostlyai_engine-2.0.0}/mostlyai/engine/training.py +0 -0
@@ -0,0 +1,319 @@
1
+ Metadata-Version: 2.4
2
+ Name: mostlyai-engine
3
+ Version: 2.0.0
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
177
+ argn.impute(data_with_missings)
178
+ ```
179
+
180
+ ### Predictions / Classification
181
+
182
+ Predict any categorical target column:
183
+
184
+ ```python
185
+ from sklearn.metrics import accuracy_score, roc_auc_score
186
+
187
+ # predict class labels for a categorical
188
+ predictions = argn.predict(data_test, target="income", n_draws=10, agg_fn="mode")
189
+
190
+ # predict class probabilities for a categorical
191
+ probabilities = argn.predict_proba(data_test, target="income", n_draws=10)
192
+
193
+ # evaluate performance
194
+ accuracy = accuracy_score(data_test["income"], predictions)
195
+ auc = roc_auc_score(data_test["income"], probabilities[:, 1])
196
+ print(f"Accuracy: {accuracy:.3f}, AUC: {auc:.3f}")
197
+ ```
198
+
199
+ ### Predictions / Regression
200
+
201
+ Predict any numerical target column:
202
+
203
+ ```python
204
+ from sklearn.metrics import mean_absolute_error
205
+
206
+ # predict target values
207
+ predictions = argn.predict(data_test, target="age", n_draws=10, agg_fn="mean")
208
+
209
+ # evaluate performance
210
+ mae = mean_absolute_error(data_test["age"], predictions)
211
+ print(f"MAE: {mae:.1f} years")
212
+ ```
213
+
214
+ ## TabularARGN for Sequential Data
215
+
216
+ For sequential data (e.g., time series or event logs), specify the context key:
217
+
218
+ ### Model Training - With Context Data
219
+
220
+ ```python
221
+ import pandas as pd
222
+ from mostlyai.engine import TabularARGN
223
+
224
+ # load sequential data
225
+ tgt_data = pd.read_csv("https://github.com/user-attachments/files/23480787/batting.csv.gz")
226
+ ctx_data = pd.read_csv("https://github.com/user-attachments/files/23480786/players.csv.gz")
227
+
228
+ # fit TabularARGN with a context key column
229
+ argn = TabularARGN(
230
+ tgt_context_key="players_id",
231
+ ctx_primary_key="id",
232
+ ctx_data=ctx_data,
233
+ max_training_time=2, # 2 minutes
234
+ verbose=0,
235
+ )
236
+ argn.fit(tgt_data)
237
+ ```
238
+
239
+ ### Sampling / Synthetic Data Generation
240
+
241
+ Generate new synthetic samples (using existing context):
242
+ ```python
243
+ argn.sample(n_samples=5)
244
+ ```
245
+
246
+ Generate new synthetic samples conditionally (using custom context and seed):
247
+
248
+ ```python
249
+ ctx_data = pd.DataFrame({
250
+ "id": ["Player1", "Player2"],
251
+ "weight": [170, 160],
252
+ "height": [70, 68],
253
+ "bats": ["R", "L"],
254
+ "throws": ["R", "L"],
255
+ })
256
+ argn.sample(ctx_data=ctx_data)
257
+ ```
258
+
259
+ ## Basic Usage of LanguageModel
260
+
261
+ 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.
262
+
263
+ **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).
264
+
265
+ ### Model Training
266
+
267
+ Load your data and train the model:
268
+
269
+ ```python
270
+ import pandas as pd
271
+ from mostlyai.engine import LanguageModel
272
+
273
+ # load data
274
+ data = pd.read_csv("https://github.com/user-attachments/files/23486562/airbnb20k.csv.gz")
275
+
276
+ # fit LanguageModel
277
+ lm = LanguageModel(
278
+ model="MOSTLY_AI/LSTMFromScratch-3m",
279
+ tgt_encoding_types={
280
+ 'neighbourhood': 'LANGUAGE_CATEGORICAL',
281
+ 'title': 'LANGUAGE_TEXT',
282
+ },
283
+ max_training_time=10, # 10 minutes
284
+ verbose=1,
285
+ )
286
+ lm.fit(data)
287
+ ```
288
+
289
+ ### Sampling / Synthetic Text Generation
290
+
291
+ Generate new synthetic samples using the trained language model:
292
+
293
+ ```python
294
+ # unconditional sampling
295
+ lm.sample(
296
+ n_samples=100,
297
+ sampling_temperature=0.8,
298
+ )
299
+ ```
300
+
301
+ ```python
302
+ # prepare seed
303
+ seed_data = pd.DataFrame({
304
+ "neighbourhood": ["Westminster", "Hackney"],
305
+ })
306
+
307
+ # conditional sampling with seed values
308
+ lm.sample(
309
+ seed_data=seed_data,
310
+ sampling_temperature=0.8,
311
+ )
312
+ ```
313
+
314
+ ## Further Examples
315
+
316
+ Example notebooks demonstrating various use cases are available in the `examples` directory:
317
+ - 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)
318
+ - 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)
319
+ - 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,268 @@
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
126
+ argn.impute(data_with_missings)
127
+ ```
128
+
129
+ ### Predictions / Classification
130
+
131
+ Predict any categorical target column:
132
+
133
+ ```python
134
+ from sklearn.metrics import accuracy_score, roc_auc_score
135
+
136
+ # predict class labels for a categorical
137
+ predictions = argn.predict(data_test, target="income", n_draws=10, agg_fn="mode")
138
+
139
+ # predict class probabilities for a categorical
140
+ probabilities = argn.predict_proba(data_test, target="income", n_draws=10)
141
+
142
+ # evaluate performance
143
+ accuracy = accuracy_score(data_test["income"], predictions)
144
+ auc = roc_auc_score(data_test["income"], probabilities[:, 1])
145
+ print(f"Accuracy: {accuracy:.3f}, AUC: {auc:.3f}")
146
+ ```
147
+
148
+ ### Predictions / Regression
149
+
150
+ Predict any numerical target column:
151
+
152
+ ```python
153
+ from sklearn.metrics import mean_absolute_error
154
+
155
+ # predict target values
156
+ predictions = argn.predict(data_test, target="age", n_draws=10, agg_fn="mean")
157
+
158
+ # evaluate performance
159
+ mae = mean_absolute_error(data_test["age"], predictions)
160
+ print(f"MAE: {mae:.1f} years")
161
+ ```
162
+
163
+ ## TabularARGN for Sequential Data
164
+
165
+ For sequential data (e.g., time series or event logs), specify the context key:
166
+
167
+ ### Model Training - With Context Data
168
+
169
+ ```python
170
+ import pandas as pd
171
+ from mostlyai.engine import TabularARGN
172
+
173
+ # load sequential data
174
+ tgt_data = pd.read_csv("https://github.com/user-attachments/files/23480787/batting.csv.gz")
175
+ ctx_data = pd.read_csv("https://github.com/user-attachments/files/23480786/players.csv.gz")
176
+
177
+ # fit TabularARGN with a context key column
178
+ argn = TabularARGN(
179
+ tgt_context_key="players_id",
180
+ ctx_primary_key="id",
181
+ ctx_data=ctx_data,
182
+ max_training_time=2, # 2 minutes
183
+ verbose=0,
184
+ )
185
+ argn.fit(tgt_data)
186
+ ```
187
+
188
+ ### Sampling / Synthetic Data Generation
189
+
190
+ Generate new synthetic samples (using existing context):
191
+ ```python
192
+ argn.sample(n_samples=5)
193
+ ```
194
+
195
+ Generate new synthetic samples conditionally (using custom context and seed):
196
+
197
+ ```python
198
+ ctx_data = pd.DataFrame({
199
+ "id": ["Player1", "Player2"],
200
+ "weight": [170, 160],
201
+ "height": [70, 68],
202
+ "bats": ["R", "L"],
203
+ "throws": ["R", "L"],
204
+ })
205
+ argn.sample(ctx_data=ctx_data)
206
+ ```
207
+
208
+ ## Basic Usage of LanguageModel
209
+
210
+ 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.
211
+
212
+ **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).
213
+
214
+ ### Model Training
215
+
216
+ Load your data and train the model:
217
+
218
+ ```python
219
+ import pandas as pd
220
+ from mostlyai.engine import LanguageModel
221
+
222
+ # load data
223
+ data = pd.read_csv("https://github.com/user-attachments/files/23486562/airbnb20k.csv.gz")
224
+
225
+ # fit LanguageModel
226
+ lm = LanguageModel(
227
+ model="MOSTLY_AI/LSTMFromScratch-3m",
228
+ tgt_encoding_types={
229
+ 'neighbourhood': 'LANGUAGE_CATEGORICAL',
230
+ 'title': 'LANGUAGE_TEXT',
231
+ },
232
+ max_training_time=10, # 10 minutes
233
+ verbose=1,
234
+ )
235
+ lm.fit(data)
236
+ ```
237
+
238
+ ### Sampling / Synthetic Text Generation
239
+
240
+ Generate new synthetic samples using the trained language model:
241
+
242
+ ```python
243
+ # unconditional sampling
244
+ lm.sample(
245
+ n_samples=100,
246
+ sampling_temperature=0.8,
247
+ )
248
+ ```
249
+
250
+ ```python
251
+ # prepare seed
252
+ seed_data = pd.DataFrame({
253
+ "neighbourhood": ["Westminster", "Hackney"],
254
+ })
255
+
256
+ # conditional sampling with seed values
257
+ lm.sample(
258
+ seed_data=seed_data,
259
+ sampling_temperature=0.8,
260
+ )
261
+ ```
262
+
263
+ ## Further Examples
264
+
265
+ Example notebooks demonstrating various use cases are available in the `examples` directory:
266
+ - 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)
267
+ - 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)
268
+ - 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.0"
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.0"
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.*")