fastpragma

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

Installation

Install latest from the GitHub repository:

Install from conda:

Install from pypi:

# Install latest from the GitHub repository:
# $ pip install git+https://github.com/risheekkumarb/fastpragma.git

# or from conda:
# $ conda install -c risheekkumarb fastpragma

# or from pypi:
# $ pip install fastpragma

Documentation

Documentation can be found hosted on this GitHub repository’s pages. Additionally you can find package manager specific guidelines on conda and pypi respectively.

Usage

fastpragma now covers four working layers:

  1. Data — declare profile and event sources with DataSource, fit a Tokenizer through PRAGMADataset, and write entity-sharded Parquet with write_kv
  2. Dataloading — load shards with pragma_dl, pragma_dls, or PRAGMADataLoader.from_path, with entity grouping, token budgets, packed events, and optional MLM masking
  3. Model — build the encoder-only PRAGMA architecture with PRAGMAModel or the pragma_model presets S, M, and L
  4. Training and tasks — pretrain with pragma_learner, resume and validate runs, extract entity embeddings, or build task dataloaders and classification models for fine-tuning

The end-to-end workflow has been run on MovieLens 100K and UCI Online Retail.

import polars as pl, torch, numpy as np
from fastai.data.external import untar_data, URLs
from fastcore.all import *
from fastai.torch_core import to_device, default_device

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

Data format

fastpragma accepts profile data with one row per entity and event data with many timestamped rows per entity. Both are declared with DataSource.

Each source can declare:

  • cats: categorical fields
  • conts: continuous numerical fields
  • signed_conts: continuous fields where sign is represented separately
  • texts: free-text fields, handled as low-cardinality tokens or BPE
  • lifelong: timestamp or milestone fields in profile data
  • time_col: the timestamp column for event data
  • entity_col: the entity identifier shared across sources

DataSource validates the declaration and supports construction from DataFrames or files.

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.

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
DataSource(columns=['user_id', 'movie_id', 'rating', 'timestamp'], name=events_df cats=['movie_id', 'rating'], conts=[], texts=[], time_col='timestamp')
shape: (5, 4)
┌─────────┬──────────┬────────┬─────────────────────┐
│ user_id ┆ movie_id ┆ rating ┆ timestamp           │
│ ---     ┆ ---      ┆ ---    ┆ ---                 │
│ i64     ┆ i64      ┆ i64    ┆ datetime[μs]        │
╞═════════╪══════════╪════════╪═════════════════════╡
│ 196     ┆ 242      ┆ 3      ┆ 1997-12-04 15:55:49 │
│ 186     ┆ 302      ┆ 3      ┆ 1998-04-04 19:22:22 │
│ 22      ┆ 377      ┆ 1      ┆ 1997-11-07 07:18:36 │
│ 244     ┆ 51       ┆ 2      ┆ 1997-11-27 05:02:03 │
│ 166     ┆ 346      ┆ 1      ┆ 1998-02-02 05:33:16 │
└─────────┴──────────┴────────┴─────────────────────┘
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
DataSource(columns=['user_id', 'age', 'gender', 'occupation', 'zip_code'], name=users cats=['gender', 'zip_code'], conts=['age'], texts=['occupation'], time_col=None)
shape: (5, 5)
┌─────────┬─────┬────────┬────────────┬──────────┐
│ user_id ┆ age ┆ gender ┆ occupation ┆ zip_code │
│ ---     ┆ --- ┆ ---    ┆ ---        ┆ ---      │
│ i64     ┆ i64 ┆ str    ┆ str        ┆ str      │
╞═════════╪═════╪════════╪════════════╪══════════╡
│ 1       ┆ 24  ┆ M      ┆ technician ┆ 85711    │
│ 2       ┆ 53  ┆ F      ┆ other      ┆ 94043    │
│ 3       ┆ 23  ┆ M      ┆ writer     ┆ 32067    │
│ 4       ┆ 24  ┆ M      ┆ technician ┆ 43537    │
│ 5       ┆ 33  ┆ F      ┆ other      ┆ 15213    │
└─────────┴─────┴────────┴────────────┴──────────┘

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. show_summary() gives an overview of sources and vocabularies.

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

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)
n_keys, n_vals
Keys: 11, Vals: 2514, BPE: none
tokenizing profile
100.00% [1/1 00:00<00:00] 
tokenizing event source 0: events_df
combining sources
100.00% [4/4 00:01<00:00] 
(11, 2514)

Dataloaders

pragma_dl creates a PyTorch DataLoader from tokenized Parquet shards. It groups records by entity, packs variable-length events under a token budget, and can apply MLM masking.

pragma_dls creates fastai train and validation DataLoaders. PRAGMADataLoader.from_path reloads a saved tokenizer and shard directory. preflight, validate_shard, validate_shards, and validate_split check shard structure, dtypes, and vocabulary sizes before loading.

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

For training with fastai, split the shard list into training and validation shards, then use preflight to validate shard structure, dtypes, and tokenizer vocabulary sizes before constructing pragma_dls.

valid_shards,train_shards = shards[-1:],shards[:-1]
preflight(train_shards, valid_shards, tok)

dls = pragma_dls(train_shards, valid_shards, tok, max_tokens=30000, valid_batches=2)

Model

PRAGMAModel is an encoder-only PRAGMA-style model with profile, event, and history encoders. It supports static and lifelong profile fields, packed variable-length events, shared key/value embeddings, RoPE, calendar features, an [USR] profile token, and an [EVT] event token.

The model includes an MLM head tied to the value vocabulary. pragma_model constructs preset sizes S, M, and L. Task-specific classification models are available through get_classification_model for fine-tuning.

model = pragma_model("S", n_keys=tok.n_keys, n_vals=tok.n_vals)

Pretraining

pragma_learner builds a PRAGMAModel and wraps it in a fastai Learner for masked-token pretraining. The workflow has been executed end to end on MovieLens 100K and UCI Online Retail, including categorical, continuous, and BPE text fields.

Implemented training support includes:

learn = pragma_learner(dls, tok.n_keys, tok.n_vals, sz='S')
if torch.cuda.is_available(): learn = learn.to_fp16()
learn.fit(1, lr=1e-3, cbs=[ThroughputCB(20)])
epoch train_loss valid_loss time
0 6.593944 6.100394 00:03 
b,_ = first(dls.valid)
b = to_device(b, default_device())
embs = entity_embs(learn.model.model, b)
len(embs), first(embs.items())[0], first(embs.items())[1].shape
(115, 4, torch.Size([192]))

Example 2: UCI Online Retail

This example uses real transaction data rather than a small teaching dataset. Each row is a purchase event, and CustomerID identifies the entity whose history will be modeled.

The preparation step removes rows without a customer, cancelled invoices, and non-positive quantities or prices. The remaining columns map directly onto the DataSource API: StockCode and Country are categorical fields, Quantity and UnitPrice are continuous fields, Description is text for BPE tokenization, and InvoiceDate supplies event time.

url = 'https://archive.ics.uci.edu/ml/machine-learning-databases/00352/Online%20Retail.xlsx'
retail = pl.read_excel(url).drop_nulls('CustomerID').with_columns(
    pl.col('CustomerID').cast(pl.Int64).alias('customer_id'),
    pl.col('InvoiceNo').cast(pl.Utf8),
    pl.col('StockCode').cast(pl.Utf8),
    pl.col('Description').fill_null('').cast(pl.Utf8)
).filter(
    ~pl.col('InvoiceNo').str.starts_with('C'),
    pl.col('Quantity') > 0,
    pl.col('UnitPrice') > 0
)

retail.head()
shape: (5, 9)
InvoiceNo StockCode Description Quantity InvoiceDate UnitPrice CustomerID Country customer_id
str str str i64 datetime[ms] f64 i64 str i64
"536365" "85123A" "WHITE HANGING HEART T-LIGHT HO… 6 2010-12-01 08:26:00 2.55 17850 "United Kingdom" 17850
"536365" "71053" "WHITE METAL LANTERN" 6 2010-12-01 08:26:00 3.39 17850 "United Kingdom" 17850
"536365" "84406B" "CREAM CUPID HEARTS COAT HANGER" 8 2010-12-01 08:26:00 2.75 17850 "United Kingdom" 17850
"536365" "84029G" "KNITTED UNION FLAG HOT WATER B… 6 2010-12-01 08:26:00 3.39 17850 "United Kingdom" 17850
"536365" "84029E" "RED WOOLLY HOTTIE WHITE HEART." 6 2010-12-01 08:26:00 3.39 17850 "United Kingdom" 17850 
purchases = DataSource(
    retail.lazy(), entity_col='customer_id',
    cats=['StockCode','Country'], conts=['Quantity','UnitPrice'],
    texts=['Description'], time_col='InvoiceDate', name='purchases'
)

purchases
DataSource(columns=['InvoiceNo', 'StockCode', 'Description', 'Quantity', 'InvoiceDate', 'UnitPrice', 'CustomerID', 'Country', 'customer_id'], name=purchases cats=['StockCode', 'Country'], conts=['Quantity', 'UnitPrice'], texts=['Description'], time_col='InvoiceDate')
shape: (5, 9)
┌───────────┬───────────┬───────────┬──────────┬───┬───────────┬───────────┬───────────┬───────────┐
│ InvoiceNo ┆ StockCode ┆ Descripti ┆ Quantity ┆ … ┆ UnitPrice ┆ CustomerI ┆ Country   ┆ customer_ │
│ ---       ┆ ---       ┆ on        ┆ ---      ┆   ┆ ---       ┆ D         ┆ ---       ┆ id        │
│ str       ┆ str       ┆ ---       ┆ i64      ┆   ┆ f64       ┆ ---       ┆ str       ┆ ---       │
│           ┆           ┆ str       ┆          ┆   ┆           ┆ i64       ┆           ┆ i64       │
╞═══════════╪═══════════╪═══════════╪══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡
│ 536365    ┆ 85123A    ┆ WHITE     ┆ 6        ┆ … ┆ 2.55      ┆ 17850     ┆ United    ┆ 17850     │
│           ┆           ┆ HANGING   ┆          ┆   ┆           ┆           ┆ Kingdom   ┆           │
│           ┆           ┆ HEART     ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ T-LIGHT   ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ HO…       ┆          ┆   ┆           ┆           ┆           ┆           │
│ 536365    ┆ 71053     ┆ WHITE     ┆ 6        ┆ … ┆ 3.39      ┆ 17850     ┆ United    ┆ 17850     │
│           ┆           ┆ METAL     ┆          ┆   ┆           ┆           ┆ Kingdom   ┆           │
│           ┆           ┆ LANTERN   ┆          ┆   ┆           ┆           ┆           ┆           │
│ 536365    ┆ 84406B    ┆ CREAM     ┆ 8        ┆ … ┆ 2.75      ┆ 17850     ┆ United    ┆ 17850     │
│           ┆           ┆ CUPID     ┆          ┆   ┆           ┆           ┆ Kingdom   ┆           │
│           ┆           ┆ HEARTS    ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ COAT      ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ HANGER    ┆          ┆   ┆           ┆           ┆           ┆           │
│ 536365    ┆ 84029G    ┆ KNITTED   ┆ 6        ┆ … ┆ 3.39      ┆ 17850     ┆ United    ┆ 17850     │
│           ┆           ┆ UNION     ┆          ┆   ┆           ┆           ┆ Kingdom   ┆           │
│           ┆           ┆ FLAG HOT  ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ WATER B…  ┆          ┆   ┆           ┆           ┆           ┆           │
│ 536365    ┆ 84029E    ┆ RED       ┆ 6        ┆ … ┆ 3.39      ┆ 17850     ┆ United    ┆ 17850     │
│           ┆           ┆ WOOLLY    ┆          ┆   ┆           ┆           ┆ Kingdom   ┆           │
│           ┆           ┆ HOTTIE    ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ WHITE     ┆          ┆   ┆           ┆           ┆           ┆           │
│           ┆           ┆ HEART.    ┆          ┆   ┆           ┆           ┆           ┆           │
└───────────┴───────────┴───────────┴──────────┴───┴───────────┴───────────┴───────────┴───────────┘
customers_df = retail.group_by('customer_id').agg(pl.col('Country').last()).lazy()
customers = DataSource(customers_df, entity_col='customer_id', cats=['Country'], name='customers', is_profile=True)
dataset = PRAGMADataset(profile=customers, events=[purchases], entity_col='customer_id', out_path='retail_data')

The profile source contains one row per customer and carries the customer’s country as static context. The purchase source remains an event stream with many timestamped rows per customer.

Passing both sources to PRAGMADataset lets fastpragma combine static customer information with the chronological purchase history before tokenization and sharding.

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

# Write tokenized entity shards.
shard_dir,n_keys,n_vals = dataset.write_kv(eval_time="2011-12-01T00:00:00", n_shards=4)



Keys: 10, Vals: 3716, BPE: 259
tokenizing profile
100.00% [1/1 00:00<00:00] 
tokenizing event source 0: purchases
combining sources
100.00% [4/4 01:21<00:00]

Here bpe_vocab_size=100 enables subword tokenization for the free-text descriptions. The resulting tokenizer reports BPE: 259, confirming that the text column produced a BPE vocabulary rather than being treated only as a categorical field.

write_kv then writes the tokenized records to four entity shards. The later dataloader, model, and learner cells use those same shards, so this is a complete data-to-pretraining test on mixed categorical, continuous, and BPE inputs.

shards = sorted(Path(shard_dir).glob('shard_*.parquet'))
valid_shards,train_shards = shards[-1:],shards[:-1]

preflight(train_shards, valid_shards, tok)
max_profile_toks,max_lifelong_toks,max_events,max_event_toks = 100,100,520,24 ## additional protection for less memory systems
dls = pragma_dls(train_shards, valid_shards, tok, max_tokens=3000, valid_batches=1, prefetch=0)

len(shards),tok.n_keys,tok.n_vals
(4, 10, 3975)
learn = pragma_learner(dls, tok.n_keys, tok.n_vals, sz='S').to_fp16()
learn.fit_one_cycle(1, lr_max=1e-3)
/usr/local/lib/python3.12/dist-packages/fastai/callback/fp16.py:47: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  self.autocast,self.learn.scaler,self.scales = autocast(dtype=dtype),GradScaler(**self.kwargs),L()
/usr/local/lib/python3.12/dist-packages/fastai/callback/fp16.py:47: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.
  self.autocast,self.learn.scaler,self.scales = autocast(dtype=dtype),GradScaler(**self.kwargs),L()
epoch train_loss valid_loss time
0 3.150082 3.353010 02:03 
learn.save('pretrain_model.pth')
Path('models/pretrain_model.pth.pth')

finetuning example

labels = retail.group_by('customer_id').agg((pl.col('Quantity') * pl.col('UnitPrice')).sum().alias('target'))
threshold = labels.select(pl.col('target').median()).item()
labels = labels.with_columns((pl.col('target') >= threshold).cast(pl.Int64).alias('label'))

rng = np.random.default_rng(42)
ids = rng.permutation(labels['customer_id'].to_numpy())
n_train = int(len(ids) * .8)
train_ids,valid_ids = ids[:n_train],ids[n_train:]
def split_shards(shards, labels, entity_col, out):
    out = Path(out)
    out.mkdir(parents=True, exist_ok=True)
    ids = labels[entity_col].to_list()
    for i,p in enumerate(shards):
        df = pl.read_parquet(p).filter(pl.col(entity_col).is_in(ids))
        if len(df): df.write_parquet(out/f'shard_{i}.parquet')
    return sorted(out.glob('shard_*.parquet'))
train_shards = split_shards(shards, train_labels, 'customer_id', 'retail_train')
valid_shards = split_shards(shards, valid_labels, 'customer_id', 'retail_valid')
retail_dls = pragma_task_dls(
    train_shards, valid_shards, train_labels, valid_labels,
    'customer_id', target_col='label', dtype=torch.long,
    max_tokens=12000, tok=tok)
from fastai.learner import Learner
from fastai.losses import CrossEntropyLossFlat
from fastai.metrics import accuracy
from fastai.optimizer import Adam
cls_model = get_classification_model(
    tok.n_keys, tok.n_vals, sz='S', n_classes=2,
    pretrain=True, pretrain_path='models/pretrain_model.pth.pth'
)

retail_learn = Learner(
    retail_dls, cls_model,
    loss_func=CrossEntropyLossFlat(),
    opt_func=Adam,
    metrics=accuracy,
    splitter=pragma_task_splitter
)
retail_learn.fit_one_cycle(1, lr_max=1e-3)
epoch train_loss valid_loss accuracy time
0 0.399607 0.398961 0.829493 01:32

Current implemented API

Data

  • DataSource wraps a polars LazyFrame, validates declared columns, and supports from_df and from_file
  • Source fields support categorical, continuous, signed continuous, textual, event-time, and lifelong/profile data
  • Tokenizer builds shared key/value vocabularies, numerical buckets, low-cardinality text tokens, and BPE vocabularies; it supports save and load
  • PRAGMADataset combines one profile source with event sources, fits or loads a tokenizer, writes entity-sharded Parquet with write_kv, and reports summaries with show_summary

Dataloading and validation

Model and pretraining

  • PRAGMAModel provides profile, event, and history encoders with shared embeddings, RoPE, calendar features, packed events, and an MLM head
  • pragma_model provides model presets S, M, and L
  • pragma_learner creates a fastai learner for masked-token pretraining
  • Checkpoint, resume, gradient-accumulation, periodic-save, and throughput callbacks are implemented
  • entity_embs extracts entity representations

Fine-tuning

Verified workflows

MovieLens 100K and UCI Online Retail both run through data preparation, tokenization, sharding, dataloading, model construction, and pretraining. Online Retail also runs through customer-level label creation, shard splitting, classification fine-tuning, and pretrained-weight loading.

Remaining work

The core data-to-pretraining pipeline and a classification fine-tuning path are implemented and tested on MovieLens 100K and UCI Online Retail.

The main remaining additions are:

  • A friendlier high-level API, including possible SourceSchema and PRAGMADataset.dataloaders() conveniences
  • A top-level model-loading API and broader save/load helpers for learners, heads, tokenizers, and weights
  • Minimal examples that run without downloading a dataset
  • Richer batch inspection and show_batch() displays
  • Higher-level embedding extraction methods
  • Additional task heads for regression, recommendation, and retrieval
  • LoRA fine-tuning and linear probing for frozen embeddings
  • Optional richer text-encoder integrations
  • More complete temporal-feature documentation
  • Additional tests and smoke-test notebooks