fastpragma

Dataloader

Assumptions: - fastpragma.data already exports DataSource, Tokenizer, and PRAGMADataset. - We use MovieLens 100k as a small stand-in for banking event histories. - User rows provide profile/static state, and rating rows provide timestamped events.

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

Assumptions: - users_src is profile state because it has one row per user and no event timestamp. - ratings_src is the event stream because each row is timestamped. - eval_time is fixed for this toy dataset so logsec features are deterministic. - Tokenized shards plus tokenizer.json are saved under out.

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 = 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')]

Reload saved tokenized data

Assumptions: - The dataloader should consume saved parquet shards, not the original raw MovieLens files. - Each shard contains rows from multiple users, with profile rows marked by event_idx == -1 and event rows by event_idx >= 0. - We first inspect one shard and one user before building batching logic.

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(shards[0]).filter(pl.col('user_id') == pl.col('user_id').first())
sample

shape: (112, 15)

user_id	key	key_id	value	val_id	val_pos	vtype	time	logsec	source_idx	event_idx	hour	dow	dom	shard_id
i64	str	i64	str	i64	i64	str	i64	f64	i32	i64	f32	f32	f32	u64
3	"age"	5	"23"	2501	0	"num"	0	0.0	-1	-1	0.0	0.0	0.0	0
3	"gender"	6	"M"	2469	0	"cat"	0	0.0	-1	-1	0.0	0.0	0.0	0
3	"occupation"	8	"writer"	2603	0	"text"	0	0.0	-1	-1	0.0	0.0	0.0	0
3	"zip_code"	10	"32067"	1203	0	"cat"	0	0.0	-1	-1	0.0	0.0	0.0	0
3	"movie_id"	7	"300"	1162	0	"cat"	889236939	33.858346	0	0	2.0	6.0	7.0	0
…	…	…	…	…	…	…	…	…	…	…	…	…	…	…
3	"rating"	9	"5"	1507	0	"cat"	889237482	0.0	0	51	2.0	6.0	7.0	0
3	"movie_id"	7	"317"	1196	0	"cat"	889237482	0.0	0	52	2.0	6.0	7.0	0
3	"rating"	9	"2"	965	0	"cat"	889237482	0.0	0	52	2.0	6.0	7.0	0
3	"movie_id"	7	"181"	932	0	"cat"	889237482	0.0	0	53	2.0	6.0	7.0	0
3	"rating"	9	"4"	1329	0	"cat"	889237482	0.0	0	53	2.0	6.0	7.0	0

Tensor helpers and limits

Assumptions: - Each key-value token is represented as (key_id, val_id, val_pos). - Profile and lifelong tensors can be padded per batch. - Event tokens will be packed flat with offsets instead of padded to a large rectangle. - Cutoffs mirror the PRAGMA paper defaults where practical: profile tokens, lifelong tokens, events per user, and tokens per event.

Convert one user into a record

Assumptions: - Static profile rows have event_idx < 0 and time == 0. - Lifelong rows have event_idx < 0 and non-zero time. - Event rows are sorted by event_idx, truncated to the latest max_events, and each event is capped at max_event_toks. - Calendar features are stored once per event.

Pack a batch

Assumptions: - Profiles and lifelong rows are small enough to pad within a batch. - Event field tokens are packed into one flat tensor. - event_offsets maps flat tokens to events. - event_user maps events to users. - history_offsets reserves one history position per user for [USR], plus one per event.

Entity ids
- entity_col is preserved as uids in each batch.
- uids contains the original entity ids, e.g. user ids.
- The model usually should not embed or consume uids.
- Use uids for debugging, evaluation joins, predictions, and traceability.
- The model should use event_user and history_offsets for batch-local entity structure.

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:str='user_id', 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=2000), 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([13, 4, 3]),
 'profile_mask': torch.Size([13, 4]),
 'profile_time': torch.Size([13, 4]),
 'lifelong': torch.Size([13, 0, 3]),
 'lifelong_mask': torch.Size([13, 0]),
 'lifelong_time': torch.Size([13, 0]),
 'event_tokens': torch.Size([1900, 3]),
 'event_offsets': torch.Size([951]),
 'event_user': torch.Size([950]),
 'event_time': torch.Size([950]),
 'cal': torch.Size([950, 3]),
 'history_offsets': torch.Size([14]),
 'uids': 13}

Masking, splits, and dataloader wrapper

Assumptions: - MLM labels reconstruct value IDs only, using event_tokens[:,1]. - Non-masked labels use -100 so PyTorch loss can ignore them. - Batches are built by token budget rather than fixed user count. - Multi-worker loading splits shard paths across workers to avoid duplicate data.

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:str='user_id', 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.2738),
 tensor([[   7, 1131,    0],
         [   9, 1507,    0],
         [   7,    1,    0],
         [   9,    1,    0],
         [   7, 1098,    0]]),
 tensor([-100, -100, 1944, 1507, -100]))

source

PRAGMADataLoader.from_path

def from_path(
    path, entity_col:str='user_id', max_tokens:int=12000, shuffle:bool=True, 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.2765), torch.Size([48970, 3]), torch.Size([48970]))

source

batch_shapes

def batch_shapes(
    b
):

Call self as a function.

Dataloader contract

The model receives padded profile/lifelong tensors and packed event tensors.

Special tokens are added by the model, not the dataloader: - profile encoder prepends [USR] - event encoder prepends [EVT] to each event - history encoder receives [USR] + [EVT] event representations

event_offsets maps packed event tokens to events. event_user maps events to batch users. history_offsets maps history tokens to users.

MLM labels reconstruct event value IDs only. Non-masked labels are -100.

Dataloader

Build and tokenize the dataset

Reload saved tokenized data

Tensor helpers and limits

Convert one user into a record

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