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:

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:

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

def prepend_batch_tok(
    x, tok
):

Call self as a function.

---

source

pack_to_padded

def pack_to_padded(
    x, offsets
):

Call self as a function.

---

source

lens_from_offsets

def lens_from_offsets(
    offsets
):

Call self as a function.

---

source

bget

def bget(
    b, k
):

Call self as a function.

---

source

KVEmb

def KVEmb(
    n_keys, n_vals, d_model, max_val_pos:int=32
):

Embed key, value, and value-position triples

---

source

sincos_pos

def sincos_pos(
    n, d
):

Call self as a function.

Rope

---

source

rope_apply

def rope_apply(
    q, k, rope_emb, pos
):

Call self as a function.

---

source

RotaryEmb

def RotaryEmb(
    dim, base:int=10000
):

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

---

source

rope_apply_complex

def rope_apply_complex(
    x, freqs
):

Call self as a function.

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__

---

source

has_flash_attn

def has_flash_attn():

Call self as a function.

---

source

can_flash_attn

def can_flash_attn(
    x
):

Call self as a function.

---

source

cu_seqlens

def cu_seqlens(
    offsets
):

Call self as a function.

---

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__

---

source

CalEmb

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

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

---

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__

---

source

add_evt_toks

def add_evt_toks(
    x, offsets, tok
):

Call self as a function.

---

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__

---

source

pad_history

def pad_history(
    z_usr, z_evt, event_time, history_offsets
):

Build padded user histories and history times

---

source

token_user_ids

def token_user_ids(
    event_offsets, event_user
):

Call self as a function.

---

source

token_event_ids

def token_event_ids(
    event_offsets
):

Call self as a function.

---

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__

---

source

MLMHead

def MLMHead(
    d_model, n_vals
):

Predict masked event value tokens

---

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.7705, grad_fn=<NllLossBackward0>)

---

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())

15687328

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)

n_keys,n_vals = len(tok.key_vocab),len(tok.val_vocab)
n_keys,n_vals

(11, 2604)

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

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=<NllLossBackward0>))

using fastai

---

source

fastai_pragma_dl

def fastai_pragma_dl(
    *args, **kwargs
):

Call self as a function.

---

source

FastaiPRAGMADL

def FastaiPRAGMADL(
    dl
):

Wrap PRAGMA batches as fastai input-target pairs

---

source

pragma_loss

def pragma_loss(
    logits, b
):

Call self as a function.

---

source

PRAGMAWrapper

def PRAGMAWrapper(
    model
):

Return MLM logits

---

source

pragma_learner

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

def pragma_dls(
    shards, tok, max_tokens:int=1500
):

Call self as a function.

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=<NllLossBackward0>)

class ShortEpochCB(Callback):
    def __init__(self, n=3): self.n = n
    def after_batch(self):
        if self.iter+1 >= self.n: raise CancelEpochException

learn = pragma_learner(dls, n_keys, n_vals)
learn.fit(1, lr=1e-3, cbs=ShortEpochCB(3))

[0, '00:23']