Metadata-Version: 2.4
Name: fastpragma
Version: 0.0.6
Summary: A practical PRAGMA model implementation for tokenized tabular event histories
Author-email: Risheek kumar B <b.risheekkumar@gmail.com>
License: Apache-2.0
Project-URL: Repository, https://github.com/risheekkumarb/fastpragma
Project-URL: Documentation, https://risheekkumarb.github.io/fastpragma/
Keywords: nbdev,pragma,deep-learning,transformers,tabular
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: polars>=1.0
Requires-Dist: numpy>=1.26
Requires-Dist: pandas>=2.0
Requires-Dist: torch>=2.2
Requires-Dist: fastai>=2.7
Requires-Dist: fastcore>=1.5
Requires-Dist: tokenizers>=0.15
Requires-Dist: fastprogress>=1.0
Requires-Dist: pyarrow>=15.0
Dynamic: license-file

# fastpragma

> An easy-to-use API for foundation model development, based on the [PRAGMA](https://arxiv.org/abs/2604.08649) framework.

## Installation

Install latest from the GitHub [repository][repo]:

```sh
pip install git+https://github.com/risheekkumarb/fastpragma.git
```

Install from [conda][conda]:

```sh
conda install -c risheekkumarb fastpragma
```

Install from [pypi][pypi]:

```sh
pip install fastpragma
```

## Documentation

Documentation can be found on this GitHub [repository][repo]'s [pages][docs]. Package-manager-specific pages are also available on [conda][conda] and [pypi][pypi].

[repo]: https://github.com/risheekkumarb/fastpragma
[docs]: https://risheekkumarb.github.io/fastpragma/
[pypi]: https://pypi.org/project/fastpragma/
[conda]: https://anaconda.org/risheekkumarb/fastpragma

## Usage

fastpragma currently exposes three main layers:

1. **Data** — declare profile/event sources with `DataSource`, then tokenize them with `PRAGMADataset`
2. **Dataloading** — read tokenized parquet shards with `pragma_dl` / `pragma_dls`
3. **Model + training** — build a PRAGMA-style model with `PRAGMAModel`, `pragma_model`, or `pragma_learner`

The API is still evolving, so this README focuses on the pieces implemented in the notebooks today.

```python
import polars as pl
from fastai.data.external import untar_data, URLs
from fastcore.all import *

from fastpragma.data import DataSource, PRAGMADataset
from fastpragma.dataloader import pragma_dl
from fastpragma.model import pragma_dls, pragma_model, pragma_learner
```

## Data format

fastpragma expects data in two broad forms:

1. **Profile data** — one row per entity, containing relatively static attributes.
2. **Event data** — many rows per entity, each with a timestamp.

Each source declares which columns are:

- `cats`: categorical fields
- `conts`: continuous numerical fields
- `signed_conts`: continuous numerical fields where sign matters separately
- `texts`: free-text fields
- `lifelong`: timestamp/milestone fields in profile data

## Example: MovieLens 100K

This example loads the classic MovieLens 100K dataset into polars DataFrames.

### Creating data sources

Use `DataSource` to declare how each DataFrame should be interpreted.

Profile sources use `is_profile=True` and normally do not need a `time_col`.
Event sources provide a `time_col`.

```python
path = untar_data(URLs.ML_100k)
events_df = pl.scan_csv(
    path/'u.data',
    separator='\t',
    has_header=False,
    new_columns=['user_id','movie_id','rating','timestamp'],
)
events_df = events_df.with_columns(pl.from_epoch('timestamp', time_unit='s').alias('timestamp'))
ratings = DataSource(events_df, entity_col='user_id', cats=['movie_id','rating'], time_col='timestamp', name='events_df')
ratings
```

```python
profile_df = pl.scan_csv(
    path/'u.user',
    separator='|',
    has_header=False,
    new_columns=['user_id','age','gender','occupation','zip_code'],
)
profile = DataSource(
    profile_df,
    entity_col='user_id',
    cats=['gender','zip_code'],
    conts=['age'],
    texts=['occupation'],
    name='users',
    is_profile=True,
)
profile
```

## Building a PRAGMADataset

`PRAGMADataset` combines one optional profile source and one or more event sources.

It fits a tokenizer, converts sources into key-value-time tokens, and writes entity-sharded parquet files.

```python
dataset = PRAGMADataset(profile=profile, events=[ratings], entity_col="user_id", out_path="data")

# Fit vocabularies and numerical buckets.
tok = dataset.fit_tokenizer(num_buckets=10, cardinality_threshold=100)

# Write tokenized entity shards.
shard_dir = dataset.write_kv(eval_time="1998-04-01T00:00:00", n_shards=4)
```

## Dataloaders

`pragma_dl` creates a PyTorch/fastai-compatible dataloader from tokenized parquet shards.

For masked-language-model pre-training, pass `mask=True` and the fitted tokenizer.

```python
shards = sorted(Path(shard_dir).glob("shard_*.parquet"))
dl = pragma_dl(shards, entity_col="user_id", max_tokens=1500, shuffle=True, tok=tok, mask=True)
```

For training with fastai, `pragma_dls` creates train/validation `DataLoaders` from the shard list.

```python
dls = pragma_dls(shards, tok=tok, max_tokens=1500)
```

## Model and training

The implemented model is an encoder-only PRAGMA-style architecture with three main pieces:

1. A profile encoder
2. An event encoder
3. A history encoder

The model predicts masked event value tokens during pre-training.

```python
# Small preset model.
model = pragma_model("S", n_keys=len(tok.key_vocab), n_vals=len(tok.val_vocab))
model
```

For quick experiments, use `pragma_learner`, which builds a compact `PRAGMAModel` and wraps it in a fastai `Learner`.

```python
learn = pragma_learner(dls, n_keys=len(tok.key_vocab), n_vals=len(tok.val_vocab))
# learn.fit(1)
```

## Current implemented API

### Data

- `DataSource`
  - wraps a polars `LazyFrame`
  - validates declared columns
  - supports `from_df(...)` and `from_file(...)`
  - handles categorical, continuous, signed continuous, textual, event-time, and lifelong/profile fields

- `Tokenizer`
  - builds key/value vocabularies
  - bucketizes numerical fields
  - tokenizes text fields
  - converts sources into key-value-time form

- `PRAGMADataset`
  - combines profile and event sources
  - fits/saves a tokenizer
  - writes sharded parquet token data with `write_kv(...)`

### Dataloading

- `PRAGMADataLoader`
  - streams parquet shards
  - groups rows by entity
  - packs variable-length records up to `max_tokens`
  - optionally applies MLM masking

- `pragma_dl`
  - convenience function returning a `DataLoader`

- `pragma_dls`
  - convenience function returning fastai `DataLoaders`

### Model

- `PRAGMAModel`
  - profile encoder
  - event encoder
  - calendar/time embeddings
  - history encoder
  - MLM head

- `pragma_model`
  - creates preset model sizes: `"S"`, `"M"`, `"L"`

- `pragma_learner`
  - creates a fastai learner for masked-token pre-training

## Planned additions

The following pieces are planned for future versions:

- A more polished public API, possibly including `SourceSchema` as a friendlier alias or replacement for `DataSource`
- A `.dataloaders()` convenience method directly on `PRAGMADataset`
- A top-level `PRAGMA.load(size="S"|"M"|"L")` model-loading API
- Better README examples using tiny synthetic data that can run without downloading MovieLens
- A richer `show_batch()` display for inspecting tokenized profile and event records
- Embedding extraction APIs such as `model.embed(dataset)` and `model.embed_record(record)`
- Task-specific heads for classification, regression, recommendation, and retrieval
- LoRA fine-tuning utilities for adapting the backbone efficiently
- Linear probing helpers for evaluating frozen embeddings
- Save/load helpers for trained learners, heads, tokenizers, and model weights
- Optional text encoder integration for richer free-text fields
- More complete documentation of temporal features, calendar features, and lifelong events
- More tests and smoke-test notebooks covering data → tokenizer → shards → dataloader → model → learner
