fastpragma

Dataloader

The dataloader consumes saved PRAGMA Parquet shards produced by PRAGMADataset.write_kv; it does not refit or modify the tokenizer. It groups rows by entity, separates profile/lifelong state from events, applies load-time limits, packs variable-length events, and emits batches for PyTorch or fastai.

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

DataSource(columns=['user_id', 'movie_id', 'rating', 'timestamp'], name=ratings 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 │
└─────────┴──────────┴────────┴─────────────────────┘

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

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    │
└─────────┴─────┴────────┴────────────┴──────────┘

Build and tokenize the dataset

PRAGMADataset creates the saved preprocessing contract: tokenizer.json plus entity-sharded Parquet. write_kv fits the tokenizer if needed, tokenizes profile and event sources, and writes the shards. The dataloader starts only after this step.

out = Path('data/ml100k_pragma')
ds = PRAGMADataset(profile=users_src, events=[ratings_src], entity_col='user_id', out_path=out)
ds.show_summary()

shape: (2, 14)

i	name	kind	rows	entities	entity_col	time_col	cats	conts	signed_conts	texts	lifelong	min_time	max_time
i64	str	str	i64	i64	str	str	i64	i64	i64	i64	i64	i64	i64
0	"users"	"profile"	943	943	"user_id"	null	2	1	0	1	0	null	null
1	"ratings"	"event"	100000	943	"user_id"	"timestamp"	2	0	0	0	0	874724710	893286638

out.mkdir(exist_ok=True)
tok,k,v = ds.write_kv(eval_time='2026-01-01', n_shards=5)
print('\n\nOutput files:\n', out.ls())

Keys: 11, Vals: 2604, BPE: none
tokenizing profile
█ |----------------------------------------| 0.00% [0/1 00:00<?]tokenizing event source 0: ratings
 |████████████████████████████████████████| 100.00% [1/1 00:00<00:00]combining sources
█ |----------------------------------------| 0.00% [0/5 00:00<?] |████████--------------------------------| 20.00% [1/5 00:00<00:01] |████████████████------------------------| 40.00% [2/5 00:00<00:01] |████████████████████████----------------| 60.00% [3/5 00:01<00:00] |████████████████████████████████--------| 80.00% [4/5 00:01<00:00] |████████████████████████████████████████| 100.00% [5/5 00:01<00:00]

Output files:
 [Path('data/ml100k_pragma/shard_1.parquet'), Path('data/ml100k_pragma/shard_2.parquet'), Path('data/ml100k_pragma/tokenizer.json'), Path('data/ml100k_pragma/shard_3.parquet'), Path('data/ml100k_pragma/shard_0.parquet'), Path('data/ml100k_pragma/shard_4.parquet')]

source

infer_entity_col

def infer_entity_col(
    path, entity_col:NoneType=None
):

Call self as a function.

Reload saved tokenized data

Reload with Tokenizer.load(out/'tokenizer.json') and discover shards with sorted(out.glob('shard_*.parquet')). Each row is one token. event_idx=-1 identifies profile/lifelong rows; non-negative event_idx identifies event rows. The loader operates entirely on these saved artifacts.

tok = Tokenizer.load(out/'tokenizer.json')
shard = pl.read_parquet(out/'shard_0.parquet')
shard.shape, shard.columns, shard.head()

((45256, 15),
 ['user_id',
  'key',
  'key_id',
  'value',
  'val_id',
  'val_pos',
  'vtype',
  'time',
  'logsec',
  'source_idx',
  'event_idx',
  'hour',
  'dow',
  'dom',
  'shard_id'],
 shape: (5, 15)
 ┌─────────┬────────────┬────────┬────────┬───┬──────┬─────┬─────┬──────────┐
 │ user_id ┆ key        ┆ key_id ┆ value  ┆ … ┆ hour ┆ dow ┆ dom ┆ shard_id │
 │ ---     ┆ ---        ┆ ---    ┆ ---    ┆   ┆ ---  ┆ --- ┆ --- ┆ ---      │
 │ i64     ┆ str        ┆ i64    ┆ str    ┆   ┆ f32  ┆ f32 ┆ f32 ┆ u64      │
 ╞═════════╪════════════╪════════╪════════╪═══╪══════╪═════╪═════╪══════════╡
 │ 3       ┆ age        ┆ 5      ┆ 23     ┆ … ┆ 0.0  ┆ 0.0 ┆ 0.0 ┆ 0        │
 │ 3       ┆ gender     ┆ 6      ┆ M      ┆ … ┆ 0.0  ┆ 0.0 ┆ 0.0 ┆ 0        │
 │ 3       ┆ occupation ┆ 8      ┆ writer ┆ … ┆ 0.0  ┆ 0.0 ┆ 0.0 ┆ 0        │
 │ 3       ┆ zip_code   ┆ 10     ┆ 32067  ┆ … ┆ 0.0  ┆ 0.0 ┆ 0.0 ┆ 0        │
 │ 3       ┆ movie_id   ┆ 7      ┆ 300    ┆ … ┆ 2.0  ┆ 6.0 ┆ 7.0 ┆ 0        │
 └─────────┴────────────┴────────┴────────┴───┴──────┴─────┴─────┴──────────┘)

shards = sorted(out.glob('shard_*.parquet'))
len(shards), shards[:2]

(5,
 [Path('data/ml100k_pragma/shard_0.parquet'),
  Path('data/ml100k_pragma/shard_1.parquet')])

sample = pl.read_parquet(out / 'shard_0.parquet').head(10)

Tensor helpers and limits

The loader represents each token as (key_id, val_id, val_pos) with a separate logsec value. Profile and lifelong sequences are padded per batch. Event tokens remain packed flat, with offsets and user indices preserving event boundaries. Limits are applied during loading, not tokenization.

Convert one entity into a record

user_rec converts one entity’s rows into profile tokens, lifelong tokens, event-token tensors, event times, and calendar features. prep_shard applies the independent limits for profile tokens, lifelong tokens, retained events, and tokens per event before record construction.

source

prep_shard

def prep_shard(
    df, entity_col, max_profile_toks:int=200, max_lifelong_toks:int=200, max_events:int=6500, max_event_toks:int=24
):

Call self as a function.

source

user_rec

def user_rec(
    df, entity_col, max_profile_toks:int=200, max_lifelong_toks:int=200, max_events:int=6500, max_event_toks:int=24
):

Call self as a function.

Pack a batch

pack_batch pads profile and lifelong tensors, concatenates event tokens into one flat tensor, and returns event_offsets, event_user, event_time, cal, history_offsets, and uids. uids are retained for joins and traceability; model structure is represented by the batch-local offset/index tensors.

detailed plan

The current fitting convention is:

PRAGMADataset.write_kv(...) is the preprocessing/fitting stage.
It fits or loads the tokenizer, tokenizes profile and event sources, assigns each entity to a shard, and writes saved parquet shards plus tokenizer.json.
The dataloader does not refit anything. It only consumes the saved files.

The saved output directory looks like:

data/ml100k_pragma/
  tokenizer.json
  shard_0.parquet
  shard_1.parquet
  ...

Each parquet shard uses this row-level convention:

one row = one key/value token
entity_col, here user_id, identifies the user/record
key_id, val_id, val_pos are the token ids used by the model
source_idx == -1 means profile state
event_idx == -1 means non-event/profile rows
event_idx >= 0 means event rows
static profile rows have event_idx < 0 and time == 0
lifelong/profile-time rows have event_idx < 0 and time != 0
event rows have source_idx >= 0, event_idx >= 0, logsec, hour, dow, dom

For one user, user_rec converts rows into:

dict(
    uid=uid,
    profile=pt,
    profile_time=ptm,
    lifelong=lt,
    lifelong_time=ltm,
    events=ets,
    event_time=event_time,
    cal=cal)

Where:

profile: static profile tokens, shape (n_profile_toks, 3)
profile_time: log-time per static profile token
lifelong: lifelong/profile-time tokens, shape (n_lifelong_toks, 3)
lifelong_time: log-time per lifelong token
events: list of event token tensors, each shape (n_event_toks, 3)
event_time: one logsec value per event
cal: one (hour,dow,dom) vector per event

For a batch, pack_batch uses this convention:

profile and lifelong tokens are padded per batch
event tokens are packed flat, not padded
event_offsets maps packed tokens back to event boundaries
event_user maps each event to its user index in the batch
history_offsets reserves one history position per user plus one per event
uids keeps the original entity ids

So overall:

raw sources -> PRAGMADataset.write_kv -> tokenizer.json + shard_*.parquet -> PRAGMADataLoader -> packed model batch

source

PRAGMADataLoader

def PRAGMADataLoader(
    shards, entity_col:NoneType=None, max_tokens:int=12000, shuffle:bool=False, tok:NoneType=None, mask:bool=False
):

Iterable dataset over saved PRAGMA parquet shards, yielding packed user batches.

source

PRAGMADataLoader.len

def __len__():

Call self as a function.

dl = DataLoader(PRAGMADataLoader(shards, max_tokens=20000), batch_size=None)
b = first(dl)
{k:(v.shape if hasattr(v,'shape') else len(v)) for k,v in b.items()}

{'profile': torch.Size([86, 4, 3]),
 'profile_mask': torch.Size([86, 4]),
 'profile_time': torch.Size([86, 4]),
 'lifelong': torch.Size([86, 0, 3]),
 'lifelong_mask': torch.Size([86, 0]),
 'lifelong_time': torch.Size([86, 0]),
 'event_tokens': torch.Size([19068, 3]),
 'event_offsets': torch.Size([9535]),
 'event_user': torch.Size([9534]),
 'event_time': torch.Size([9534]),
 'cal': torch.Size([9534, 3]),
 'history_offsets': torch.Size([87]),
 'uids': 86}

Masking, splits, and dataloader wrapper

mask_batch performs event-value MLM masking using token-, event-, and user-local key-level probabilities. pragma_dl returns a PyTorch DataLoader; PRAGMADataLoader.from_path reloads a tokenizer and shard directory. The loader supports token-budget batching, optional shuffling, masking, and worker-based shard splitting.

MLM masking assumptions

We mask only event value tokens, not profile or lifelong tokens.

For each user record, selected mask positions come from the union of three sources:

Token-level masking: each event token is selected independently with probability p_tok.
Event-level masking: each event is selected independently with probability p_event, and all value tokens in that event are selected.
Semantic key-level masking: keys are selected independently per user with probability p_key, and all event values for those keys are selected for that user.

After positions are selected:

selected positions become MLM targets using their original val_id
most selected input values are replaced with [MASK]
a small fraction are replaced with [UNK]
[UNK] positions are excluded from the loss by setting their label to -100

So the effective target mask is:

This follows the paper’s intent while treating key-level masking as sample-local: selected keys affect only the user record where they were sampled, not other users in the same batch.

source

pragma_dl

def pragma_dl(
    shards, entity_col:NoneType=None, max_tokens:int=12000, shuffle:bool=False, tok:NoneType=None, mask:bool=False,
    num_workers:int=0
):

Call self as a function.

b = first(pragma_dl(shards, max_tokens=50000, shuffle=True, tok=tok, mask=True))

test_eq(b['event_labels'].shape, (len(b['event_tokens']),))
test_eq(b['mlm_mask'].shape, (len(b['event_tokens']),))
test_eq((b['event_labels'][~b['mlm_mask']] == -100).all().item(), True)
test_eq((b['event_labels'][b['mlm_mask']] != -100).all().item(), True)

b['mlm_mask'].float().mean(), b['event_tokens'][:5], b['event_labels'][:5]

(tensor(0.3047),
 tensor([[   7, 1162,    0],
         [   9,  965,    0],
         [   7, 1173,    0],
         [   9,  965,    0],
         [   7, 1227,    0]]),
 tensor([-100, -100, -100, -100, -100]))

source

PRAGMADataLoader.from_path

def from_path(
    path, entity_col:NoneType=None, max_tokens:int=12000, shuffle:bool=False, tok:NoneType=None, mask:bool=True,
    num_workers:int=0
):

Call self as a function.

dl = PRAGMADataLoader.from_path(out, max_tokens=50000)
b = first(dl)
b['mlm_mask'].float().mean(), b['event_tokens'].shape, b['event_labels'].shape

(tensor(0.3049), torch.Size([49040, 3]), torch.Size([49040]))

source

batch_shapes

def batch_shapes(
    b
):

Call self as a function.

Dataloader contract

The batch contains padded profile/lifelong tensors and packed event tensors:

profile, profile_mask, profile_time, lifelong, lifelong_mask, lifelong_time, event_tokens, event_offsets, event_user, event_time, cal, history_offsets, uids, event_labels, and mlm_mask.

The model adds [USR] and [EVT]; the dataloader supplies only token data, temporal features, structure indices, and MLM labels.

Dataloader

Build and tokenize the dataset

infer_entity_col

Reload saved tokenized data

Tensor helpers and limits

Convert one entity into a record

prep_shard

user_rec

Pack a batch

detailed plan

PRAGMADataLoader

PRAGMADataLoader.__len__

Masking, splits, and dataloader wrapper

MLM masking assumptions

pragma_dl

PRAGMADataLoader.from_path

batch_shapes

Dataloader contract

PRAGMADataLoader.len