## Dataloader Assumptions: - `fastpragma.data` already exports [`DataSource`](https://risheekkumarb.github.io/fastpragma/data.html#datasource), [`Tokenizer`](https://risheekkumarb.github.io/fastpragma/data.html#tokenizer), and [`PRAGMADataset`](https://risheekkumarb.github.io/fastpragma/data.html#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. ``` python 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 │ └─────────┴──────────┴────────┴─────────────────────┘ ``` python 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`. ``` python 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
``` python 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= 0`. - We first inspect one shard and one user before building batching logic. ``` python 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 │ └─────────┴────────────┴────────┴────────┴───┴──────┴─────┴─────┴──────────┘) ``` python 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')]) ``` python 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: ``` python 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`](https://risheekkumarb.github.io/fastpragma/dataloader.html#user_rec) converts rows into: ``` python 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`](https://risheekkumarb.github.io/fastpragma/dataloader.html#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: ``` python raw sources -> PRAGMADataset.write_kv -> tokenizer.json + shard_*.parquet -> PRAGMADataLoader -> packed model batch ``` ------------------------------------------------------------------------ source ### PRAGMADataLoader ``` python 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\_\_ ``` python def __len__(): ``` *Call self as a function.* ``` python 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 ``` python 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.* ``` python 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 ``` python 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.* ``` python 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 ``` python 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`.