fastpragma

An easy-to-use API for foundation-model development based on the PRAGMA framework.

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

Installation

pip install fastpragma

Install the latest development version from GitHub:

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

Conda installation:

conda install -c risheekkumarb fastpragma

Documentation

• Documentation
• GitHub repository
• PyPI
• Conda

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

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 LazyFrames. 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

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.

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:

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.

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:

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:

'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.

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:

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:

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.

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.