Model

This notebook implements the current PRAGMA encoder over the saved dataloader contract. The model consumes padded profile/lifelong tensors and packed event tokens; it never refits the tokenizer or consumes raw source tables.

The pipeline is ProfileEncoder → EventEncoder → HistoryEncoder → MLMHead. It uses shared key/value embeddings, model-owned [USR] and [EVT] tokens, calendar embeddings, event-time features, and masked event-value prediction.

helpers

Rope

Attention variants

---

source

RopeMHA

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__

VarlenMHA handles the packed event representation. event_tokens is flat and event_offsets marks each event’s span, so attention stays within each event. FlashAttention is used when available on a compatible CUDA tensor; otherwise the implementation pads temporarily with PyTorch attention and returns the flat result.

---

source

VarlenMHA

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

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

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__

CalEmb converts each event’s (hour, dow, dom) values into a model-width vector using periodic sine/cosine features followed by an MLP. This gives the event encoder cyclical calendar information rather than treating hour or weekday as ordinary distances.

---

source

CalEmb

def CalEmb(
    d_model, n_cal:int=3
):

Same as nn.Module, but no need for subclasses to call super().__init__

Profile Encoder

ProfileEncoder embeds padded profile and optional lifelong tokens, prepends a learnable [USR] token, applies the profile transformer, and returns the [USR] representation as one vector per entity.

---

source

ProfileEncoder

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__

Event Encoder

EventEncoder embeds flat event tokens, prepends [EVT] within each event, applies variable-length attention using event_offsets, and returns both one summary vector per event and contextualized token representations for MLM.

---

source

EventEncoder

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__

History Encoder

HistoryEncoder combines the profile representation with event representations for each entity in the sequence [USR, event_1, event_2, ...]. event_user reconstructs entity membership, while event_time supplies temporal positions. It returns updated entity and event representations.

---

source

HistoryEncoder

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__

MLMHead predicts value IDs for selected event tokens from three contexts: token representation, event representation, and entity representation. Its logits are compared with the dataloader’s event_labels only at positions selected by mlm_mask.

---

source

MLMHead

def MLMHead(
    d_model
):

Predict masked event value tokens

PRAGMAModel runs the full representation pipeline and returns h_usr, h_evt, and MLM logits. When labels are present it also returns the selected labels and cross-entropy loss. pragma_model('S'|'M'|'L', n_keys, n_vals) supplies the tested model presets.

---

source

PRAGMAModel

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__

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.5428, grad_fn=<AddBackward0>)

---

source

pragma_model

def pragma_model(
    sz, n_keys, n_vals, p:float=0.1
):

Call self as a function.

n_keys,n_vals = 60,28_000
m = pragma_model('S', n_keys, n_vals)
sum(p.numel() for p in m.parameters())

10283328

Example

out = Path('data/ml100k_pragma')
shards = sorted(out.glob('shard_*.parquet'))
len(shards),shards[0]

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

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)

def mlm_loss(out): return nn.functional.cross_entropy(out['logits'], out['labels'])

tiny_dl = pragma_dl(shards, max_tokens=15000, shuffle=False, mask=True, tok=tok)
b = next(iter(tiny_dl))
m = PRAGMAModel(tok.n_keys, tok.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([64, 32]),
 torch.Size([7341, 32]),
 torch.Size([5117, 2604]),
 tensor(8.4505, grad_fn=<NllLossBackward0>))