## 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`](https://risheekkumarb.github.io/fastpragma/finetune.html#pragmataskloader) wraps [`PRAGMADataLoader`](https://risheekkumarb.github.io/fastpragma/dataloader.html#pragmadataloader), looks up the target for each batch `uid`, and yields `(batch, target)` pairs. [`pragma_task_dl`](https://risheekkumarb.github.io/fastpragma/finetune.html#pragma_task_dl) and [`pragma_task_dls`](https://risheekkumarb.github.io/fastpragma/finetune.html#pragma_task_dls) provide the single- and train/validation-loader APIs. ------------------------------------------------------------------------ source ### pragma_task_dls ``` python 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 ``` python 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 ``` python 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`](https://risheekkumarb.github.io/fastpragma/finetune.html#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 ``` python 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`](https://risheekkumarb.github.io/fastpragma/finetune.html#pragma_user_emb) reuses the backbone’s profile, event, calendar, and history encoders to produce one vector per entity. [`PRAGMATaskModel`](https://risheekkumarb.github.io/fastpragma/finetune.html#pragmataskmodel) applies a task-specific [`TaskHead`](https://risheekkumarb.github.io/fastpragma/finetune.html#taskhead) to those vectors. ------------------------------------------------------------------------ source ### PRAGMATaskModel ``` python def PRAGMATaskModel( backbone, head ): ``` *Same as `nn.Module`, but no need for subclasses to call `super().__init__`* ------------------------------------------------------------------------ source ### pragma_user_emb ``` python def pragma_user_emb( m, b ): ``` *Call self as a function.* [`pragma_model`](https://risheekkumarb.github.io/fastpragma/model.html#pragma_model) supplies the pretrained backbone and its `d_model` width. [`get_classification_model`](https://risheekkumarb.github.io/fastpragma/finetune.html#get_classification_model) and [`get_regression_model`](https://risheekkumarb.github.io/fastpragma/finetune.html#get_regression_model) construct task models, optionally loading pretrained weights through [`load_pretrained`](https://risheekkumarb.github.io/fastpragma/finetune.html#load_pretrained). [`pragma_task_splitter`](https://risheekkumarb.github.io/fastpragma/finetune.html#pragma_task_splitter) exposes separate backbone and head parameter groups to fastai. ------------------------------------------------------------------------ source ### pragma_task_splitter ``` python def pragma_task_splitter( m ): ``` *Call self as a function.* ### loading model and using it ------------------------------------------------------------------------ source ### get_model ``` python 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 ``` python def load_pretrained( m, path ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### get_regression_model ``` python 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 ``` python 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.* ``` python 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. ``` python 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. ``` python 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**: ``` python 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. ``` python 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. ``` python 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. ``` python 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. ``` python 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`](https://risheekkumarb.github.io/fastpragma/model.html#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. ``` python 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} ``` python 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') ``` python pretrain_path = Path('models/pragma_pretrain.pth') pretrain_path.exists() ``` True ``` python 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