Metadata-Version: 2.4
Name: premixdb
Version: 0.1.3
Summary: Data mixtures for LLM pretraining that you can reproduce and trace
Requires-Python: >=3.12
Description-Content-Type: text/markdown
Requires-Dist: protobuf<7,>=6.33.5
Requires-Dist: blake3>=1.0.10
Requires-Dist: packaging>=24
Requires-Dist: ipython<10,>=9.17.1
Requires-Dist: tokenizers<0.23.1,>=0.22; sys_platform == "darwin" and platform_machine == "x86_64"
Requires-Dist: tokenizers<0.24,>=0.23.1; sys_platform != "darwin" or platform_machine != "x86_64"
Requires-Dist: datatrove==0.3.0; sys_platform == "darwin" and platform_machine == "x86_64"
Requires-Dist: datatrove==0.10.1; sys_platform != "darwin" or platform_machine != "x86_64"
Requires-Dist: numpy<2,>=1.26; sys_platform == "darwin" and platform_machine == "x86_64"
Requires-Dist: fasteners>=0.19
Requires-Dist: fasttext-wheel==0.9.2; sys_platform == "darwin" and platform_machine == "x86_64"
Requires-Dist: fasttext-numpy2-wheel>=0.9.2; sys_platform != "darwin" or platform_machine != "x86_64"
Requires-Dist: regex>=2025.10.22
Requires-Dist: tldextract<6,>=5.3
Requires-Dist: spacy<4,>=3.8
Requires-Dist: fsspec[http]>=2023.12.0
Requires-Dist: marin-dupekit==0.1.1
Requires-Dist: pyarrow<24,>=18
Requires-Dist: pyroaring<2,>=1.2
Requires-Dist: transformers==4.57.6; sys_platform == "darwin" and platform_machine == "x86_64"
Requires-Dist: transformers<6,>=5.10; sys_platform != "darwin" or platform_machine != "x86_64"
Requires-Dist: sentence-transformers<6,>=5.7
Requires-Dist: torch<2.3,>=2.2; sys_platform == "darwin" and platform_machine == "x86_64"
Requires-Dist: torch<3,>=2.5; sys_platform != "darwin" or platform_machine != "x86_64"
Requires-Dist: einops<1,>=0.8
Requires-Dist: datasets<6,>=5.0.1
Provides-Extra: huggingface

# premixdb

premixdb is a declarative library for building reproducible data mixtures to improve llm pretraining.

## Quick start

Run a premixdb shell without installing the library:

```bash
uvx --python 3.12 premixdb shell
```

The shell has already run these commands so you have `p` and `db` available:

```python
import premixdb as p
db = p.PremixDB()
```

You can also use the tiny_shakespeare demo corpus of his plays:

```python
from torch.utils.data import DataLoader

dataset = db.Corpus('demo').query().mix()[0].train.torch()
next(iter(DataLoader(dataset, batch_size=1)))
```

```
{'input_ids': tensor([[45472, 10426,  1677,  ...,   475,   286,   477]]),
 'attention_mask': tensor([[1, 1, 1,  ..., 1, 1, 1]]),
 'labels': tensor([[45472, 10426,  1677,  ...,   475,   286,   477]])}
```

## Full example

Run from the repository root using the small checked-in C4 and S2ORC-derived
samples. Capture snapshots, filter them, and mix the results for PyTorch training.
Language and quality models download on first use.

```python
from pathlib import Path
from torch.utils.data import DataLoader

c4 = db.Corpus(
    "c4",
    p.Source.read_jsonl(Path("examples/data/c4.jsonl"), limit=8),
)
papers = db.Corpus(
    "papers",
    p.Source.read_jsonl(Path("examples/data/s2orc-train.jsonl"), key_column="id"),
)
query = c4.union(papers).query(
    steps=[
        p.where(p.text.characters >= 200),
        p.dedupe(),
        p.where(p.language.en >= 0.75),
        p.where(p.quality.educational_value >= 1.0),
    ]
)

mixtures = query.mix(tokens=256, replacement=True, sequence_length=64)

dataset = mixtures[0]
print(dataset.preview())
batch = next(iter(DataLoader(dataset.train.torch(), batch_size=1)))
```

## Installation

To add as a dependency to your project use:

```bash
uv add premixdb
```

- [Runnable examples](examples/README.md)
- [RegMix training search](examples/11_c4_pretraining_ablation.py): train C4 topic/content-type mixtures and select weights by validation loss
- [Model data recipes](examples/recipes/README.md): T5/C4, Falcon, Gopher, LLaMA 1, GPT-3 with local sample inputs
- [Curation](docs/curation.md) · [Fields](docs/enrichment.md) · [Training](docs/training.md)
- [Storage](docs/persistence.md) · [Internals](docs/architecture.md) · [Development](docs/development.md)
