## Model
This notebook implements a PRAGMA-style encoder that consumes the
tokenized batches from `fastpragma.dataloader`. The goal here is a clear
working version before adding production details such as RoPE, optimized
varlen attention, LoRA, or larger model configs.
Key assumptions to keep in mind for the model:
- The dataloader consumes **already-tokenized parquet shards**, not raw
source tables.
- `PRAGMADataset.write_kv(...)` is the preprocessing stage: it
fits/loads the tokenizer, tokenizes sources, assigns shards, and
writes `tokenizer.json` plus `shard_*.parquet`.
- The model should not refit or modify the tokenizer.
- Each parquet row is one key/value token.
- Token identity is represented by `(key_id, val_id, val_pos)`.
- `key_id` and `val_id` come from shared saved vocabularies.
- `val_pos` is the within-field value/subword position, not a global
sequence position.
- Profile state and event history are separated by row metadata:
- `source_idx == -1` and `event_idx == -1` means profile/profile-like
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`.
- Static profile tokens are padded per batch as `profile`.
- Lifelong/profile-time tokens are padded per batch as `lifelong`.
- Event tokens are **packed flat**, not padded into a
`(batch, events, tokens)` tensor.
- `event_offsets` gives start/end positions for each event in
`event_tokens`.
- `event_user` maps each packed event to a batch-local user index.
- `history_offsets` maps each user to its history span: one `[USR]`
position plus one position per event.
- Special tokens are added by the **model**, not the dataloader:
- profile encoder prepends `[USR]`.
- event encoder prepends `[EVT]` to each event.
- history encoder receives the user representation plus event
representations.
- Calendar features are stored once per event as
`cal = (hour, dow, dom)`.
- `event_time` stores one log-time value per event.
- `profile_time` and `lifelong_time` store one log-time value per
profile/lifelong token.
- Static profile time is normally zero.
- Lifelong rows use log-time distance from the evaluation point.
- Event rows use log-time distance from the entity’s latest event.
- The model should use batch-local structure fields like
`event_offsets`, `event_user`, and `history_offsets`.
- `uids` are preserved for debugging, prediction joins, evaluation, and
traceability.
- The model generally should **not** embed or consume `uids`.
- Truncation assumptions mirror PRAGMA-style limits:
- `max_profile_toks = 200`
- `max_lifelong_toks = 64`
- `max_events = 6500`
- `max_event_toks = 24`
- `max_tokens = 12000` per dataloader batch in this notebook
- For each user, events are sorted and the most recent `max_events` are
kept.
- Each event is capped to `max_event_toks`.
- Profile tokens are capped separately from lifelong tokens.
- Batch size is dynamic by token budget, not fixed number of users.
- MLM masking applies only to **event value tokens**, not profile or
lifelong tokens.
- MLM labels reconstruct `event_tokens[:,1]`, i.e. original `val_id`.
- Non-masked labels are `-100` for PyTorch loss ignoring.
- Masking combines token-level, event-level, and per-user key-level
masking.
- Some selected positions become `[UNK]` and are excluded from the loss.
- The model’s MLM head should predict value IDs for event tokens.
- The expected model contract is:
``` python
batch_fields = 'profile profile_mask profile_time lifelong lifelong_mask lifelong_time event_tokens event_offsets event_user event_time cal history_offsets uids event_labels mlm_mask'.split()
```
- Typical first batch shapes look like:
``` python
dict(profile=(64, 4, 3), profile_mask=(64, 4), profile_time=(64, 4), lifelong=(64, 0, 3), lifelong_mask=(64, 0), lifelong_time=(64, 0), event_tokens=(11518, 3), event_offsets=(5760,), event_user=(5759,), event_time=(5759,), cal=(5759, 3), history_offsets=(65,), uids=64, event_labels=(11518,), mlm_mask=(11518,))
```
Most importantly: the model should be designed around **padded
profile/lifelong sequences + packed event tokens + offset tensors**,
rather than assuming a dense rectangular event tensor.
### helpers
------------------------------------------------------------------------
source
### prepend_batch_tok
``` python
def prepend_batch_tok(
x, tok
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### pack_to_padded
``` python
def pack_to_padded(
x, offsets
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### lens_from_offsets
``` python
def lens_from_offsets(
offsets
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### bget
``` python
def bget(
b, k
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### KVEmb
``` python
def KVEmb(
n_keys, n_vals, d_model, max_val_pos:int=32
):
```
*Embed key, value, and value-position triples*
------------------------------------------------------------------------
source
### sincos_pos
``` python
def sincos_pos(
n, d
):
```
*Call self as a function.*
### Rope
------------------------------------------------------------------------
source
### rope_apply
``` python
def rope_apply(
q, k, rope_emb, pos
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### RotaryEmb
``` python
def RotaryEmb(
dim, base:int=10000
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### rope_apply_complex
``` python
def rope_apply_complex(
x, freqs
):
```
*Call self as a function.*
### Attention variants
------------------------------------------------------------------------
source
### RopeMHA
``` python
def RopeMHA(
d_model, n_heads:int=4, p:float=0.1, rope:bool=False
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### has_flash_attn
``` python
def has_flash_attn():
```
*Call self as a function.*
------------------------------------------------------------------------
source
### can_flash_attn
``` python
def can_flash_attn(
x
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### cu_seqlens
``` python
def cu_seqlens(
offsets
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### VarlenMHA
``` python
def VarlenMHA(
d_model, n_heads:int=4, p:float=0.1
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
### making encoder layers
------------------------------------------------------------------------
source
### Encoder
``` python
def Encoder(
d_model, n_heads:int=4, n_layers:int=2, d_ff:NoneType=None, p:float=0.1, rope:bool=False
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### EncoderLayer
``` python
def EncoderLayer(
d_model, n_heads:int=4, d_ff:NoneType=None, p:float=0.1, rope:bool=False
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### CalEmb
``` python
def CalEmb(
d_model, n_cal:int=3
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### ProfileEncoder
``` python
def ProfileEncoder(
emb, d_model, n_heads:int=4, n_layers:int=2, p:float=0.1
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### add_evt_toks
``` python
def add_evt_toks(
x, offsets, tok
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### EventEncoder
``` python
def EventEncoder(
emb, n_vals, d_model, n_heads:int=4, n_layers:int=2, p:float=0.1
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### pad_history
``` python
def pad_history(
z_usr, z_evt, event_time, history_offsets
):
```
*Build padded user histories and history times*
------------------------------------------------------------------------
source
### token_user_ids
``` python
def token_user_ids(
event_offsets, event_user
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### token_event_ids
``` python
def token_event_ids(
event_offsets
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### HistoryEncoder
``` python
def HistoryEncoder(
d_model, n_heads:int=4, n_layers:int=2, p:float=0.1
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### MLMHead
``` python
def MLMHead(
d_model, n_vals
):
```
*Predict masked event value tokens*
------------------------------------------------------------------------
source
### PRAGMAModel
``` python
def PRAGMAModel(
n_keys, n_vals, d_model:int=128, n_heads:int=4, prof_layers:int=1, event_layers:int=2, hist_layers:int=2,
p:float=0.1
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
``` python
bs,n_keys,n_vals,d = 2,10,30,16
b = dict(
profile=torch.randint(0, 10, (bs, 3, 3)),
profile_mask=torch.ones(bs, 3).bool(),
profile_time=torch.zeros(bs, 3),
lifelong=torch.empty(bs, 0, 3).long(),
lifelong_mask=torch.empty(bs, 0).bool(),
lifelong_time=torch.empty(bs, 0),
event_tokens=torch.randint(0, 10, (5, 3)),
event_offsets=torch.tensor([0, 2, 5]),
event_user=torch.tensor([0, 1]),
event_time=torch.tensor([1., 2.]),
cal=torch.tensor([[12, 2, 15], [18, 5, 20]]),
history_offsets=torch.tensor([0, 1, 2]),
uids=['u1', 'u2'],
event_labels=torch.randint(0, n_vals, (5,)),
mlm_mask=torch.tensor([1, 0, 1, 0, 1]).bool())
m = PRAGMAModel(n_keys, n_vals, d_model=d, n_heads=4, prof_layers=1, event_layers=1, hist_layers=1, p=0.)
res = m(b)
test_eq(res['h_usr'].shape, (bs, d))
test_eq(res['h_evt'].shape, (2, d))
test_eq(res['logits'].shape, (3, n_vals))
test_eq(res['labels'].shape, (3,))
test_eq(res['loss'].ndim, 0)
torch.isfinite(res['loss'])
res['loss']
```
tensor(3.7705, grad_fn=)
------------------------------------------------------------------------
source
### pragma_model
``` python
def pragma_model(
sz, n_keys, n_vals, p:float=0.1
):
```
*Call self as a function.*
``` python
n_keys,n_vals = 60,28_000
m = pragma_model('S', n_keys, n_vals)
sum(p.numel() for p in m.parameters())
```
15687328
## Example
``` python
out = Path('data/ml100k_pragma')
shards = sorted(out.glob('shard_*.parquet'))
len(shards),shards[0]
```
(5, Path('data/ml100k_pragma/shard_0.parquet'))
``` python
tok = Tokenizer.load(out/'tokenizer.json')
dl = pragma_dl(shards, max_tokens=12000, shuffle=False, mask=True, tok=tok)
b = next(iter(dl))
test_eq(set('profile profile_mask profile_time lifelong lifelong_mask lifelong_time event_tokens event_offsets event_user event_time cal history_offsets uids event_labels mlm_mask'.split()) <= set(b), True)
```
``` python
n_keys,n_vals = len(tok.key_vocab),len(tok.val_vocab)
n_keys,n_vals
```
(11, 2604)
``` python
def mlm_loss(out): return nn.functional.cross_entropy(out['logits'], out['labels'])
```
``` python
tiny_dl = pragma_dl(shards, max_tokens=1500, shuffle=False, mask=True, tok=tok)
b = next(iter(tiny_dl))
m = PRAGMAModel(n_keys, n_vals, d_model=32, n_heads=2, prof_layers=1, event_layers=1, hist_layers=1, p=0.)
outp = m(b)
loss = mlm_loss(outp)
loss.backward()
outp['h_usr'].shape,outp['h_evt'].shape,outp['logits'].shape,loss
```
(torch.Size([7, 32]),
torch.Size([712, 32]),
torch.Size([358, 2604]),
tensor(8.5259, grad_fn=))
### using fastai
------------------------------------------------------------------------
source
### fastai_pragma_dl
``` python
def fastai_pragma_dl(
*args, **kwargs
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### FastaiPRAGMADL
``` python
def FastaiPRAGMADL(
dl
):
```
*Wrap PRAGMA batches as fastai input-target pairs*
------------------------------------------------------------------------
source
### pragma_loss
``` python
def pragma_loss(
logits, b
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### PRAGMAWrapper
``` python
def PRAGMAWrapper(
model
):
```
*Return MLM logits*
------------------------------------------------------------------------
source
### pragma_learner
``` python
def pragma_learner(
dls, n_keys, n_vals, d_model:int=32, n_heads:int=2, opt_func:function=Adam
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### pragma_dls
``` python
def pragma_dls(
shards, tok, max_tokens:int=1500
):
```
*Call self as a function.*
``` python
dls = pragma_dls(shards, tok)
learn = pragma_learner(dls, n_keys, n_vals)
xb,yb = next(iter(dls.train))
loss = learn.loss_func(learn.model(xb), yb)
test_eq(loss.ndim, 0)
loss
```
tensor(8.0833, grad_fn=)
``` python
class ShortEpochCB(Callback):
def __init__(self, n=3): self.n = n
def after_batch(self):
if self.iter+1 >= self.n: raise CancelEpochException
```
``` python
learn = pragma_learner(dls, n_keys, n_vals)
learn.fit(1, lr=1e-3, cbs=ShortEpochCB(3))
```
[0, '00:23']