Metadata-Version: 2.4
Name: premixdb
Version: 0.1.1
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.24,>=0.23.1
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: transformers<6,>=5.10
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
Requires-Dist: datasets<6,>=5.0.1; 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:

```bash
import premixdb as p
db = PremixDB()
```

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

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

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

```
{'input_ids': tensor([[   35,    52,  7336,  ...,   628, 50256,  5097]]),
 'attention_mask': tensor([[1, 1, 1,  ..., 1, 1, 1]]),
 'labels': tensor([[   35,    52,  7336,  ...,   628, 50256,  5097]])}
```

## Full example

Capture a corpus snapshot. Filter the snapshot with a query. Mix query results into datasets. Train with Pytorch.

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

c4 = db.corpus(
    "c4",
    p.HuggingFaceSource("datablations/c4-filter-small"),
    limit=8,
)
oscar = db.corpus(
    "oscar",
    p.HuggingFaceSource("datablations/oscar-filter-small"),
    limit=8,
)
query = c4.union(oscar).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, sequence_length=64)

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

## Installation

To add as a dependency to your project use:

```bash
uv add premixdb
```

- [Runnable examples](examples/README.md)
- [Curation](docs/curation.md) · [Fields](docs/enrichment.md) · [Training](docs/training.md)
- [Storage](docs/persistence.md) · [Internals](docs/architecture.md) · [Development](docs/development.md)
