Finetune

The fine-tuning layer attaches task heads to the pretrained PRAGMA entity representation. The implemented path supports labelled entity batches, classification, regression, pretrained-weight loading, frozen-backbone training followed by full fine-tuning, and fastai metrics.

Labels remain in a Polars table keyed by entity_col. PRAGMATaskLoader wraps PRAGMADataLoader, looks up the target for each batch uid, and yields (batch, target) pairs. pragma_task_dl and pragma_task_dls provide the single- and train/validation-loader APIs.

---

source

pragma_task_dls

def pragma_task_dls(
    train_shards, valid_shards, train_labels, valid_labels, entity_col, target_col:str='target',
    max_tokens:int=12000, dtype:dtype=torch.int64, tok:NoneType=None, mask:bool=False
):

Call self as a function.

---

source

pragma_task_dl

def pragma_task_dl(
    shards, labels, entity_col, target_col:str='target', max_tokens:int=12000, shuffle:bool=False,
    dtype:dtype=torch.int64, tok:NoneType=None, mask:bool=False
):

Call self as a function.

---

source

PRAGMATaskLoader

def PRAGMATaskLoader(
    shards, labels, entity_col, target_col:str='target', max_tokens:int=12000, shuffle:bool=False, tok:NoneType=None,
    mask:bool=False, dtype:dtype=torch.int64
):

Yield PRAGMA entity batches paired with labels from a keyed Polars table.

TaskHead is a normalized feed-forward prediction head over the entity vector. Set n_out for classification or use n_out=1, squeeze=True for scalar regression.

---

source

TaskHead

def TaskHead(
    d_model, n_out, hidden:NoneType=None, p:float=0.1, squeeze:bool=False
):

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

pragma_user_emb reuses the backbone’s profile, event, calendar, and history encoders to produce one vector per entity. PRAGMATaskModel applies a task-specific TaskHead to those vectors.

---

source

PRAGMATaskModel

def PRAGMATaskModel(
    backbone, head
):

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

---

source

pragma_user_emb

def pragma_user_emb(
    m, b
):

Call self as a function.

pragma_model supplies the pretrained backbone and its d_model width. get_classification_model and get_regression_model construct task models, optionally loading pretrained weights through load_pretrained. pragma_task_splitter exposes separate backbone and head parameter groups to fastai.

---

source

pragma_task_splitter

def pragma_task_splitter(
    m
):

Call self as a function.

loading model and using it

---

source

get_model

def get_model(
    keys, vals, n_out, sz:str='S', squeeze:bool=False, head:NoneType=None, pretrain:bool=False,
    pretrain_path:NoneType=None
):

Call self as a function.

---

source

load_pretrained

def load_pretrained(
    m, path
):

Call self as a function.

---

source

get_regression_model

def get_regression_model(
    keys, vals, sz:str='S', head:NoneType=None, pretrain:bool=False, pretrain_path:NoneType=None
):

Call self as a function.

---

source

get_classification_model

def get_classification_model(
    keys, vals, sz:str='S', n_classes:int=2, head:NoneType=None, pretrain:bool=False, pretrain_path:NoneType=None
):

Call self as a function.

import nbdev; nbdev.nbdev_export()

Example

Load the UCI credit dataset, create entity IDs and binary targets, infer categorical and continuous columns, make a reproducible entity-level split, and convert the data into PRAGMA sources and tokenized shards.

from ucimlrepo import fetch_ucirepo

credit = fetch_ucirepo(id=27)
features = pl.from_pandas(credit.data.features).with_row_index('entity_id')
targets = pl.from_pandas(credit.data.targets).with_row_index('entity_id').rename({'A16':'target'})
targets = targets.with_columns((pl.col('target') == '+').cast(pl.Int64))

cats = [c for c,t in features.schema.items() if t == pl.String]
conts = [c for c,t in features.schema.items() if t in (pl.Int64, pl.Float64)]

rng = np.random.default_rng(42)
ids = rng.permutation(features['entity_id'].to_numpy())
n_train = int(len(ids) * .8)
train_ids,valid_ids = ids[:n_train],ids[n_train:]

train_labels = targets.filter(pl.col('entity_id').is_in(train_ids))
valid_labels = targets.filter(pl.col('entity_id').is_in(valid_ids))

profile = DataSource(features.select('entity_id').lazy(), entity_col='entity_id', is_profile=True, name='credit_profile')
event_df = features.with_columns(pl.lit('2020-01-01').str.to_datetime().alias('event_time'))
events = DataSource(event_df.lazy(), cats=cats, conts=conts, entity_col='entity_id', time_col='event_time', name='credit_event')

out = Path('data/credit_pragma')
ds = PRAGMADataset(profile=profile, events=[events], entity_col='entity_id', out_path=out)
out.mkdir(exist_ok=True)
out_dir,k,v = ds.write_kv(eval_time='2026-01-01', n_shards=5)
shards = sorted(out_dir.glob('shard_*.parquet'))

Keys: 20, Vals: 137, BPE: none
tokenizing profile

100.00% [1/1 00:00<00:00]

tokenizing event source 0: credit_event
combining sources

100.00% [5/5 00:01<00:00]

Define a helper that filters each tokenized shard by the entity IDs in a label table and writes the filtered shards to a new directory.

def split_shards(shards, labels, entity_col, out):
    out = Path(out)
    out.mkdir(parents=True, exist_ok=True)
    ids = labels[entity_col].to_list()
    for i,p in enumerate(shards):
        df = pl.read_parquet(p).filter(pl.col(entity_col).is_in(ids))
        if len(df): df.write_parquet(out/f'shard_{i}.parquet')
    return sorted(out.glob('shard_*.parquet'))

Create separate shard directories containing only the entities assigned to training or validation. In the example, it is a random entity-level split:

ids = rng.permutation(features['entity_id'].to_numpy())
n_train = int(len(ids) * .8)
train_ids,valid_ids = ids[:n_train],ids[n_train:]

So 80% of entities go to training and 20% to validation. The fixed seed (42) makes the split reproducible.

train_shards = split_shards(shards, train_labels, 'entity_id', out/'train')
valid_shards = split_shards(shards, valid_labels, 'entity_id', out/'valid')

Example

The implemented example builds entity-level labels, creates a reproducible train/validation split, filters tokenized shards with split_shards, constructs task dataloaders, and trains classification or regression heads. The same APIs also support loading a pretrained checkpoint before fine-tuning.

cls_model = get_classification_model(k, v)

cls_dls = pragma_task_dls(train_shards, valid_shards, train_labels, valid_labels, 'entity_id',dtype=torch.long)
cls_learn = Learner(cls_dls, cls_model, loss_func=nn.CrossEntropyLoss(), opt_func=Adam, metrics=accuracy,splitter=pragma_task_splitter)

cls_learn.freeze()
cls_learn.fit_one_cycle(2, 1e-3)

cls_learn.unfreeze()
cls_learn.fit_one_cycle(3, slice(1e-5, 1e-4))

cls_preds,cls_targs = cls_learn.get_preds(dl=cls_learn.dls.valid)
accuracy(cls_preds, cls_targs)

epoch	train_loss	valid_loss	accuracy	time
0	0.693744	0.672728	0.608696	00:01
1	0.691743	0.654900	0.608696	00:00

epoch	train_loss	valid_loss	accuracy	time
0	0.683182	0.654728	0.608696	00:00
1	0.685655	0.652424	0.608696	00:00
2	0.684643	0.651335	0.608696	00:00

TensorBase(0.6087)

Extract the learned classification model’s user embeddings without tracking gradients; the final expression checks their shape.

b,y = first(cls_dls.train)
emb = get_emb(cls_model, to_device(b))
emb.shape

torch.Size([552, 192])

Prepare regression labels, split the tokenized shards using the same entity split, build regression dataloaders, train the regression head and backbone, and evaluate with RMSE.

reg_targets = (
    pl.from_pandas(credit.data.targets)
    .with_row_index('entity_id')
    .rename({'A16': 'target'})
    .with_columns(
        pl.col('target').replace({'+': 1, '-': 0}).cast(pl.Int64),
        pl.when(pl.col('entity_id').is_in(train_ids))
        .then(pl.lit('train'))
        .otherwise(pl.lit('valid'))
        .alias('split'))
)

reg_train_labels = reg_targets.filter(pl.col('split') == 'train')
reg_valid_labels = reg_targets.filter(pl.col('split') == 'valid')

reg_train_shards = split_shards(shards, reg_train_labels, 'entity_id', out/'reg_train')
reg_valid_shards = split_shards(shards, reg_valid_labels, 'entity_id', out/'reg_valid')

reg_dls = pragma_task_dls(
    reg_train_shards, reg_valid_shards,
    reg_train_labels, reg_valid_labels,
    'entity_id',
    dtype=torch.float32)

reg_model = get_regression_model(k, v)

reg_learn = Learner(reg_dls, reg_model, loss_func=MSELossFlat(), opt_func=Adam, metrics=rmse, splitter=pragma_task_splitter)

reg_learn.freeze()
reg_learn.fit_one_cycle(2, 1e-3)

reg_learn.unfreeze()
reg_learn.fit_one_cycle(3, slice(1e-5, 1e-4))

reg_preds,reg_targs = reg_learn.get_preds(dl=reg_learn.dls.valid)
rmse(reg_preds, reg_targs)

epoch	train_loss	valid_loss	_rmse	time
0	0.589223	0.471452	0.686624	00:00
1	0.575011	0.267196	0.516910	00:00

epoch	train_loss	valid_loss	_rmse	time
0	0.341010	0.265658	0.515421	00:00
1	0.341791	0.249973	0.499973	00:00
2	0.332000	0.249444	0.499444	00:00

TensorBase(0.4994)

Pretrain then finetune

The complete transfer workflow is implemented: pretrain PRAGMAModel on event-value MLM, save the learner checkpoint, load its backbone weights with get_classification_model(..., pretrain=True, pretrain_path=...), then train the task head and optionally unfreeze the backbone.

labels = (
    pl.from_pandas(credit.data.targets)
    .with_row_index('entity_id')
    .rename({'A16':'target'})
    .with_columns(pl.col('target').replace({'+':1, '-':0}).cast(pl.Int64)))

rng = np.random.default_rng(42)
ids = rng.permutation(labels['entity_id'].to_numpy())
n_train = int(len(ids) * .8)
train_ids,valid_ids = ids[:n_train],ids[n_train:]

train_labels = labels.filter(pl.col('entity_id').is_in(train_ids))
valid_labels = labels.filter(pl.col('entity_id').is_in(valid_ids))

for d in [out/'train', out/'valid']:
    d.mkdir(parents=True, exist_ok=True)
    for p in d.glob('shard_*.parquet'): p.unlink()

train_shards = split_shards(shards, train_labels, 'entity_id', out/'train')
valid_shards = split_shards(shards, valid_labels, 'entity_id', out/'valid')

dict(
    total=len(labels),
    train=len(train_labels),
    valid=len(valid_labels),
    train_positive=train_labels['target'].mean(),
    valid_positive=valid_labels['target'].mean(),
    train_shards=len(train_shards),
    valid_shards=len(valid_shards))

{'total': 690,
 'train': 552,
 'valid': 138,
 'train_positive': 0.4583333333333333,
 'valid_positive': 0.391304347826087,
 'train_shards': 5,
 'valid_shards': 5}

tok = ds.tokenizer
pretrain_dls = pragma_dls(train_shards, valid_shards, tok, max_tokens=1500, valid_batches=1, prefetch=0)
pretrain_learn = pragma_learner(pretrain_dls, tok.n_keys, tok.n_vals,sz='S')

pretrain_learn.fit_one_cycle(1, lr_max=1e-3)
pretrain_learn.save('pragma_pretrain')

epoch	train_loss	valid_loss	time
0	4.484961	3.812326	00:01

Path('models/pragma_pretrain.pth')

pretrain_path = Path('models/pragma_pretrain.pth')
pretrain_path.exists()

True

cls_dls = pragma_task_dls(train_shards, valid_shards, train_labels, valid_labels, 'entity_id',dtype=torch.long)
cls_model = get_classification_model(k,v, sz='S',n_classes=2,pretrain=True,pretrain_path=pretrain_path)
cls_learn = Learner(cls_dls, cls_model, loss_func=nn.CrossEntropyLoss(), opt_func=Adam, metrics=accuracy, splitter=pragma_task_splitter)

cls_learn.freeze()
cls_learn.fit_one_cycle(1, 1e-3)

cls_learn.unfreeze()
cls_learn.fit_one_cycle(1, slice(1e-5, 1e-4))

epoch	train_loss	valid_loss	accuracy	time
0	0.721444	0.735318	0.391304	00:00

epoch	train_loss	valid_loss	accuracy	time
0	0.717924	0.734432	0.391304	00:00