fastpragma

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

Installation

Install latest from the GitHub repository:

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

Install from conda:

conda install -c risheekkumarb fastpragma

Install from pypi:

pip install fastpragma

Documentation

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

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.

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.

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

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.

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.

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.

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.

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

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