mostlyai-engine 1.7.1__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.
- mostlyai_engine-2.0.0/PKG-INFO +319 -0
- mostlyai_engine-2.0.0/README.md +268 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/__init__.py +14 -2
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_common.py +115 -0
- mostlyai_engine-2.0.0/mostlyai/engine/_language/interface.py +322 -0
- mostlyai_engine-2.0.0/mostlyai/engine/_tabular/interface.py +657 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/logging.py +14 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/pyproject.toml +4 -2
- mostlyai_engine-1.7.1/PKG-INFO +0 -205
- mostlyai_engine-1.7.1/README.md +0 -155
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/.gitignore +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/LICENSE +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_dtypes.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/__init__.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/__init__.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/categorical.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/datetime.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/numeric.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/language/text.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/__init__.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/categorical.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/character.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/datetime.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/itt.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/lat_long.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_encoding_types/tabular/numeric.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/__init__.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/common.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/encoding.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/__init__.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/base.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/hf_engine.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/engine/vllm_engine.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/generation.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/lstm.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/tokenizer_utils.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/training.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_language/xgrammar_utils.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_memory.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/__init__.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/argn.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/common.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/encoding.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/fairness.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/generation.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_tabular/training.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_training_utils.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/_workspace.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/analysis.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/domain.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/encoding.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/generation.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/random_state.py +0 -0
- {mostlyai_engine-1.7.1 → mostlyai_engine-2.0.0}/mostlyai/engine/splitting.py +0 -0
- {mostlyai_engine-1.7.1 → 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
|
+

|
|
55
|
+
[](https://mostly-ai.github.io/mostlyai-engine/)
|
|
56
|
+
[](https://pypi.org/project/mostlyai-engine/)
|
|
57
|
+

|
|
58
|
+

|
|
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 [](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/flat.ipynb)
|
|
318
|
+
- TabularARGN for sequential data [](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)
|
|
319
|
+
- LanguageModel for textual data [](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
|
+

|
|
4
|
+
[](https://mostly-ai.github.io/mostlyai-engine/)
|
|
5
|
+
[](https://pypi.org/project/mostlyai-engine/)
|
|
6
|
+

|
|
7
|
+

|
|
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 [](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/flat.ipynb)
|
|
267
|
+
- TabularARGN for sequential data [](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)
|
|
268
|
+
- LanguageModel for textual data [](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__ = [
|
|
25
|
-
|
|
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.*")
|