Pretrain

Pretraining consumes existing tokenizer.json and Parquet shards. It does not fit the tokenizer. The fastai wrapper builds a Learner around PRAGMAModel, with token-budget dataloaders, event-value MLM, validation, checkpointing, resume helpers, gradient accumulation, throughput reporting, and entity-embedding extraction.

---

source

warn_big_shards

def warn_big_shards(
    shards, max_gb:float=0.5
):

Call self as a function.

---

source

shard_info

def shard_info(
    p, checksum:bool=False
):

Call self as a function.

---

source

fhash

def fhash(
    p, n:int=1048576
):

Call self as a function.

---

source

fsz

def fsz(
    p
):

Call self as a function.

---

source

load_state

def load_state(
    path
):

Call self as a function.

---

source

save_state

def save_state(
    path, epoch, batch, shard_i:NoneType=None, seed:NoneType=None, extra:NoneType=None
):

Call self as a function.

---

source

resume_state

def resume_state(
    shards, state_path, shuffle:bool=False
):

Call self as a function.

---

source

pragma_learner

def pragma_learner(
    dls, n_keys, n_vals, sz:str='S', opt_func:function=Adam, p:float=0.1
):

Call self as a function.

---

source

pragma_dls

def pragma_dls(
    train_shards, valid_shards, tok, max_tokens:int=1500, valid_batches:NoneType=None, prefetch:int=2, **kwargs
):

Call self as a function.

---

source

fastai_pragma_dl

def fastai_pragma_dl(
    shards, tok, max_tokens:int=1500, shuffle:bool=False, n_batches:NoneType=None, n_skip:int=0, prefetch:int=2,
    **kwargs
):

Call self as a function.

---

source

ResumeCB

def ResumeCB(
    state_path:str='models/pretrain_state.json', every_batches:int=500, every_secs:NoneType=None,
    fname:str='pretrain', seed:int=42, with_opt:bool=True
):

Basic class handling tweaks of the training loop by changing a Learner in various events

---

source

PeriodicSaveCB

def PeriodicSaveCB(
    state_path:str='models/pretrain_state.json', every_batches:int=500, every_secs:NoneType=None,
    fname:str='pretrain', seed:int=42, with_opt:bool=True
):

Basic class handling tweaks of the training loop by changing a Learner in various events

---

source

GradAccumCB

def GradAccumCB(
    n_acc:int=1
):

Basic class handling tweaks of the training loop by changing a Learner in various events

---

source

ThroughputCB

def ThroughputCB(
    every:int=50
):

Basic class handling tweaks of the training loop by changing a Learner in various events

shards = sorted(Path('data').glob('shard_*.parquet'))
tok = Tokenizer.load('data/tokenizer.json')
valid_shards,train_shards = shards[-1:],shards[:-1]

dls = pragma_dls(train_shards, valid_shards, tok, max_tokens=1500, valid_batches=1, prefetch=0)
learn = pragma_learner(dls, tok.n_keys, tok.n_vals)

learn.fit_one_cycle(1, lr_max=1e-3, cbs=[GradientClip(1.0)])

epoch	train_loss	valid_loss	time
0	5.484083	5.298973	00:08

Getting Entity Embeddings

entity_embs(model, batch) runs the model without gradients and returns a dictionary mapping each original uid to its h_usr embedding. This is the supported extraction path for downstream retrieval, clustering, evaluation joins, and task construction.

---

source

entity_embs

def entity_embs(
    model, b
):

Return uid->embedding for one batch

b,_ = first(dls.valid)
embs = entity_embs(learn.model.model, to_device(b, 'cuda') if torch.cuda.is_available() else b)
len(embs), first(embs.items())[0], first(embs.items())[1].shape

(4, 4, torch.Size([192]))

Validation before run

preflight(train_shards, valid_shards, tok) checks train/validation overlap, required shard columns, expected dtypes, tokenizer bounds, and optional shard sizes before training. validate_shard, validate_shards, and validate_split are available for narrower checks.

---

source

bad_dtypes

def bad_dtypes(
    schema,
    dtypes:dict={'key_id': Int64, 'val_id': Int64, 'val_pos': Int64, 'source_idx': Int32, 'event_idx': Int64, 'logsec': Float64, 'hour': Float32, 'dow': Float32, 'dom': Float32}
):

Call self as a function.

---

source

missing_cols

def missing_cols(
    schema, entity_col:str='user_id'
):

Call self as a function.

---

source

shard_schema

def shard_schema(
    p
):

Call self as a function.

---

source

validate_split

def validate_split(
    train_shards, valid_shards
):

Call self as a function.

---

source

preflight

def preflight(
    train_shards, valid_shards, tok, max_gb:float=0.5, entity_col:NoneType=None
):

Call self as a function.

Resume pretraining

resumed_pragma_dls rebuilds dataloaders from saved state and skips completed training batches. resumed_pragma_learner additionally constructs the learner, loads the checkpoint, restores optimizer state when requested, and supports fp16 training.

---

source

resumed_pragma_dls

def resumed_pragma_dls(
    train_shards, valid_shards, tok, state_path, max_tokens:int=1500, valid_batches:NoneType=None, seed:int=42,
    prefetch:int=2, fname:str='pretrain', **kwargs
):

Call self as a function.

pf = preflight(train_shards, valid_shards, tok)
dls,resume = resumed_pragma_dls(train_shards, valid_shards, tok, 'models/pretrain_state.json', max_tokens=20000, valid_batches=20)
pf['split'],resume

({'n_train': 3, 'n_valid': 1},
 {'epoch': 0,
  'batch': 0,
  'train_shards': [Path('data/shard_0.parquet'),
   Path('data/shard_1.parquet'),
   Path('data/shard_2.parquet')],
  'ckpt': 'pretrain_0_0',
  'state': None})

---

source

resumed_pragma_learner

def resumed_pragma_learner(
    train_shards, valid_shards, tok, state_path, n_keys, n_vals, max_tokens:int=1500, valid_batches:NoneType=None,
    seed:int=42, prefetch:int=2, fname:str='pretrain', sz:str='S', opt_func:function=Adam, p:float=0.1,
    fp16:bool=True, with_opt:bool=True, **kwargs
):

Call self as a function.

Example

from fastai.data.external import URLs, untar_data

path = untar_data(URLs.ML_100k)
events_df = pl.scan_csv(path/'u.data', separator='\t', has_header=False, new_columns=['user_id','movie_id','rating','timestamp'])
events_df = events_df.with_columns(pl.from_epoch('timestamp', time_unit='s').alias('timestamp'))
ratings = DataSource(events_df, entity_col='user_id', cats=['movie_id','rating'], time_col='timestamp', name='events_df')
ratings

DataSource(columns=['user_id', 'movie_id', 'rating', 'timestamp'], name=events_df cats=['movie_id', 'rating'], conts=[], texts=[], time_col='timestamp')
shape: (5, 4)
┌─────────┬──────────┬────────┬─────────────────────┐
│ user_id ┆ movie_id ┆ rating ┆ timestamp           │
│ ---     ┆ ---      ┆ ---    ┆ ---                 │
│ i64     ┆ i64      ┆ i64    ┆ datetime[μs]        │
╞═════════╪══════════╪════════╪═════════════════════╡
│ 196     ┆ 242      ┆ 3      ┆ 1997-12-04 15:55:49 │
│ 186     ┆ 302      ┆ 3      ┆ 1998-04-04 19:22:22 │
│ 22      ┆ 377      ┆ 1      ┆ 1997-11-07 07:18:36 │
│ 244     ┆ 51       ┆ 2      ┆ 1997-11-27 05:02:03 │
│ 166     ┆ 346      ┆ 1      ┆ 1998-02-02 05:33:16 │
└─────────┴──────────┴────────┴─────────────────────┘

profile_df = pl.scan_csv(path/'u.user', separator='|', has_header=False, new_columns=['user_id','age','gender','occupation','zip_code'])
profile = DataSource(profile_df, entity_col='user_id', cats=['gender','zip_code'], conts=['age'], texts=['occupation'], name='users', is_profile=True)
profile

DataSource(columns=['user_id', 'age', 'gender', 'occupation', 'zip_code'], name=users cats=['gender', 'zip_code'], conts=['age'], texts=['occupation'], time_col=None)
shape: (5, 5)
┌─────────┬─────┬────────┬────────────┬──────────┐
│ user_id ┆ age ┆ gender ┆ occupation ┆ zip_code │
│ ---     ┆ --- ┆ ---    ┆ ---        ┆ ---      │
│ i64     ┆ i64 ┆ str    ┆ str        ┆ str      │
╞═════════╪═════╪════════╪════════════╪══════════╡
│ 1       ┆ 24  ┆ M      ┆ technician ┆ 85711    │
│ 2       ┆ 53  ┆ F      ┆ other      ┆ 94043    │
│ 3       ┆ 23  ┆ M      ┆ writer     ┆ 32067    │
│ 4       ┆ 24  ┆ M      ┆ technician ┆ 43537    │
│ 5       ┆ 33  ┆ F      ┆ other      ┆ 15213    │
└─────────┴─────┴────────┴────────────┴──────────┘

dataset = PRAGMADataset(profile=profile, events=[ratings], entity_col="user_id", out_path="data")

# Fit vocabularies and numerical buckets.
tok = dataset.fit_tokenizer( num_buckets=10, cardinality_threshold=100)

# Write tokenized entity shards.
shard_dir,_,_ = dataset.write_kv(eval_time="1998-04-01T00:00:00", n_shards=4)

Keys: 11, Vals: 2514, BPE: none
tokenizing profile

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

tokenizing event source 0: events_df
combining sources

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

n_keys, n_vals = tok.n_keys, tok.n_vals

shards = sorted(Path(shard_dir).glob('shard_*.parquet'))
valid_shards = shards[-1:]
train_shards = shards[:-1]
len(train_shards), len(valid_shards), warn_big_shards(shards)

(3, 1, [])

dls = pragma_dls(train_shards, valid_shards, tok, max_tokens=1500, valid_batches=1, prefetch=0)
learn = pragma_learner(dls, n_keys, n_vals)
learn.fit_one_cycle(1, lr_max=1e-3)

epoch	train_loss	valid_loss	time
0	5.450467	5.350104	00:07

testing failure

class FailAfterCB(Callback):
    order = 80
    def __init__(self, n=1): store_attr()
    def before_fit(self): self.i = 0
    def after_batch(self):
        if not self.training: return
        self.i += 1
        if self.i>=self.n: raise Exception(f'failed after {self.i} train batches')

with tempfile.TemporaryDirectory() as d:
    state_path = Path(d)/'state.json'
    cb = PeriodicSaveCB(state_path, every_batches=1, fname='pretrain')
    cb.learn,cb.training,cb.epoch = FakeLearn(),True,0
    cb.before_fit()
    cb.after_batch()
    st = load_state(state_path)
    test_eq(st['epoch'], 0)
    test_eq(st['batch'], 1)
    test_eq(cb.learn.saved, [('pretrain_0_1', True)])
    cb.after_epoch()
    st = load_state(state_path)
    test_eq(st['epoch'], 1)
    test_eq(st['batch'], 0)

import nbdev; nbdev.nbdev_export()