Metadata-Version: 2.4
Name: fastpragma
Version: 0.0.20
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.42.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.

fastpragma turns profile data and timestamped entity events into tokenized, entity-sharded Parquet data for PRAGMA-style pretraining and downstream tasks.

## Installation

```bash
pip install fastpragma
```

Install the latest development version from GitHub:

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

Conda installation:

```bash
conda install -c risheekkumarb fastpragma
```

## Documentation

- [Documentation](https://risheekkumarb.github.io/fastpragma/)
- [GitHub repository](https://github.com/risheekkumarb/fastpragma)
- [PyPI](https://pypi.org/project/fastpragma/)
- [Conda](https://anaconda.org/risheekkumarb/fastpragma)

## Overview

The library is organized into four layers:

1. **Data** — declare profile and event sources with `DataSource`, fit a `Tokenizer`, and write tokenized entity shards with `PRAGMADataset`.
2. **Dataloading** — reload saved tokenizers and shards, group rows by entity, pack events under token budgets, and optionally apply MLM masking.
3. **Model** — build the encoder-only PRAGMA architecture with profile, event, and history encoders.
4. **Training and tasks** — pretrain with masked event-value prediction, extract entity embeddings, and fine-tune classification or regression heads.

The complete workflow has been exercised on MovieLens 100K and UCI Online Retail.

## Imports

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

from fastpragma.data import *
from fastpragma.dataloader import *
from fastpragma.model import *
from fastpragma.pretrain import *
from fastpragma.finetune import *
```

## Data format

fastpragma accepts two kinds of sources:

- **Profile sources**: one row per entity, declared with `is_profile=True`.
- **Event sources**: many timestamped rows per entity, declared with `time_col`.

Each `DataSource` can declare:

- `cats`: categorical fields
- `conts`: continuous numerical fields
- `signed_conts`: continuous fields whose sign is represented separately
- `texts`: text fields, handled as categorical tokens or BPE
- `lifelong`: timestamped or milestone profile fields
- `entity_col`: the shared entity identifier
- `time_col`: the event timestamp column

Sources use Polars `LazyFrame`s. `DataSource.from_df` adapts a pandas DataFrame, and `DataSource.from_file` selects the appropriate Polars scanner for common file types.

## End-to-end example

### 1. Declare profile and event sources

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

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

### 2. Fit the tokenizer and write shards

`PRAGMADataset` combines one optional profile source with one or more event sources. `fit_tokenizer` builds shared key/value vocabularies, numerical buckets, low-cardinality text tokens, and optional BPE state.

```python
dataset = PRAGMADataset(profile=profile, events=[ratings], entity_col='user_id', out_path='data/ml100k_pragma')
tok = dataset.fit_tokenizer(num_buckets=10, cardinality_threshold=100)
shard_dir,n_keys,n_vals = dataset.write_kv(eval_time='1998-04-01T00:00:00', n_shards=4)
```

The output contains `tokenizer.json` and `shard_*.parquet`. The returned tuple is `(shard_dir, n_keys, n_vals)`.

The tokenizer can also be persisted explicitly:

```python
tok.save(Path(shard_dir)/'tokenizer.json')
tok = Tokenizer.load(Path(shard_dir)/'tokenizer.json')
```

### 3. Build dataloaders

The dataloader consumes saved shards rather than raw source tables. It separates profile and lifelong state from event history, applies context limits, packs event tokens, and preserves entity IDs for joins and evaluation.

```python
shards = sorted(Path(shard_dir).glob('shard_*.parquet'))
valid_shards,train_shards = shards[-1:],shards[:-1]
preflight(train_shards, valid_shards, tok)
dls = pragma_dls(train_shards, valid_shards, tok, max_tokens=3000, valid_batches=1)
```

For a single PyTorch dataloader:

```python
dl = pragma_dl(shards, entity_col='user_id', max_tokens=3000, shuffle=True, tok=tok, mask=True)
```

For a saved tokenizer and shard directory, use `PRAGMADataLoader.from_path`.

Batches contain padded profile/lifelong tensors and packed event tensors, including:

```python
'profile profile_mask profile_time lifelong lifelong_mask lifelong_time event_tokens event_offsets event_user event_time cal history_offsets uids event_labels mlm_mask'.split()
```

### 4. Build and pretrain the model

`pragma_model` provides the tested presets `S`, `M`, and `L`. The model adds `[USR]` and `[EVT]` tokens itself, uses packed event attention, and predicts masked event values.

```python
model = pragma_model('S', n_keys=tok.n_keys, n_vals=tok.n_vals)
learn = pragma_learner(dls, tok.n_keys, tok.n_vals, sz='S')
learn.fit_one_cycle(1, lr_max=1e-3)
```

For CUDA mixed precision:

```python
if torch.cuda.is_available(): learn = learn.to_fp16()
```

The pretraining module also provides:

- `save_state`, `load_state`, and `resume_state`
- `resumed_pragma_dls` and `resumed_pragma_learner`
- `PeriodicSaveCB`, `ResumeCB`, `GradAccumCB`, and `ThroughputCB`
- `entity_embs` for extracting per-entity representations

Extract embeddings from a batch with:

```python
b,_ = first(dls.valid)
embs = entity_embs(learn.model.model, b)
```

`embs` maps each original entity ID to its model-width embedding.

## Fine-tuning

Fine-tuning uses labelled entity data and the pretrained user representation. The current task API supports classification and regression.

```python
labels = pl.DataFrame({'user_id':[1,2,3], 'label':[0,1,0]})
train_dls = pragma_task_dls(train_shards, valid_shards, labels, labels, 'user_id', target_col='label', dtype=torch.long, max_tokens=3000, tok=tok)
model = get_classification_model(tok.n_keys, tok.n_vals, sz='S', n_classes=2, pretrain=True, pretrain_path='models/pretrain_model.pth.pth')
learn = Learner(train_dls, model, loss_func=CrossEntropyLossFlat(), metrics=accuracy, splitter=pragma_task_splitter)
```

For scalar regression, use `get_regression_model` with `dtype=torch.float` and a regression loss/metric.

`load_pretrained` loads model weights from a learner or model checkpoint. `pragma_task_splitter` exposes separate backbone and head parameter groups for staged training or freezing.

## Validation and data safety

Use the validation helpers before training:

- `validate_shard`
- `validate_shards`
- `validate_split`
- `preflight`

These check shard structure, required columns, dtypes, tokenizer compatibility, and train/validation entity overlap.

## Implemented API

### Data

- `DataSource`
- `Tokenizer`
- `PRAGMADataset`
- `DataSource.from_df`
- `DataSource.from_file`
- `Tokenizer.save` and `Tokenizer.load`
- `PRAGMADataset.fit_tokenizer`
- `PRAGMADataset.write_kv`
- `PRAGMADataset.show_summary`

### Dataloading

- `PRAGMADataLoader`
- `PRAGMADataLoader.from_path`
- `pragma_dl`
- `pragma_dls`
- `preflight`
- `validate_shard`
- `validate_shards`
- `validate_split`

### Model and pretraining

- `PRAGMAModel`
- `pragma_model`
- `pragma_learner`
- `entity_embs`
- `save_state`
- `load_state`
- `resume_state`
- `resumed_pragma_dls`
- `resumed_pragma_learner`
- `PeriodicSaveCB`
- `ResumeCB`
- `GradAccumCB`
- `ThroughputCB`

### Fine-tuning

- `pragma_task_dl`
- `pragma_task_dls`
- `TaskHead`
- `PRAGMATaskModel`
- `get_classification_model`
- `get_regression_model`
- `load_pretrained`
- `pragma_task_splitter`

## Current scope

The core data-to-pretraining pipeline and classification fine-tuning path are implemented and tested on MovieLens 100K and UCI Online Retail. Further work includes friendlier high-level convenience APIs, richer batch inspection, recommendation and retrieval heads, LoRA fine-tuning, linear probing, and broader model/checkpoint helpers.

