fastpragma

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

Usage

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

# Core data API
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
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.

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)
Keys: 11, Vals: 2514, BPE: none
tokenizing profile
█

 |----------------------------------------| 0.00% [0/1 00:00<?]tokenizing event source 0: events_df

 |████████████████████████████████████████| 100.00% [1/1 00:00<00:00]combining sources
█

 |----------------------------------------| 0.00% [0/4 00:00<?]
 |██████████------------------------------| 25.00% [1/4 00:00<00:01]
 |████████████████████--------------------| 50.00% [2/4 00:00<00:00]
 |██████████████████████████████----------| 75.00% [3/4 00:01<00:00]
 |████████████████████████████████████████| 100.00% [4/4 00:01<00:00]

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))

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