## Pretrain Pretraining consumes existing `tokenizer.json` and Parquet shards. It does not fit the tokenizer. The fastai wrapper builds a `Learner` around [`PRAGMAModel`](https://risheekkumarb.github.io/fastpragma/model.html#pragmamodel), with token-budget dataloaders, event-value MLM, validation, checkpointing, resume helpers, gradient accumulation, throughput reporting, and entity-embedding extraction. ------------------------------------------------------------------------ source ### warn_big_shards ``` python def warn_big_shards( shards, max_gb:float=0.5 ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### shard_info ``` python def shard_info( p, checksum:bool=False ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### fhash ``` python def fhash( p, n:int=1048576 ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### fsz ``` python def fsz( p ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### load_state ``` python def load_state( path ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### save_state ``` python def save_state( path, epoch, batch, shard_i:NoneType=None, seed:NoneType=None, extra:NoneType=None ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### resume_state ``` python def resume_state( shards, state_path, shuffle:bool=False ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### pragma_learner ``` python 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 ``` python 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 ``` python 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 ``` python 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 ``` python 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 ``` python def GradAccumCB( n_acc:int=1 ): ``` *Basic class handling tweaks of the training loop by changing a `Learner` in various events* ------------------------------------------------------------------------ source ### ThroughputCB ``` python def ThroughputCB( every:int=50 ): ``` *Basic class handling tweaks of the training loop by changing a `Learner` in various events* ``` python 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 ``` python def entity_embs( model, b ): ``` *Return uid-\>embedding for one batch* ``` python 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`](https://risheekkumarb.github.io/fastpragma/pretrain.html#validate_shard), [`validate_shards`](https://risheekkumarb.github.io/fastpragma/pretrain.html#validate_shards), and [`validate_split`](https://risheekkumarb.github.io/fastpragma/pretrain.html#validate_split) are available for narrower checks. ------------------------------------------------------------------------ source ### bad_dtypes ``` python 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 ``` python def missing_cols( schema, entity_col:str='user_id' ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### shard_schema ``` python def shard_schema( p ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### validate_split ``` python def validate_split( train_shards, valid_shards ): ``` *Call self as a function.* ------------------------------------------------------------------------ source ### preflight ``` python 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`](https://risheekkumarb.github.io/fastpragma/pretrain.html#resumed_pragma_dls) rebuilds dataloaders from saved state and skips completed training batches. [`resumed_pragma_learner`](https://risheekkumarb.github.io/fastpragma/pretrain.html#resumed_pragma_learner) additionally constructs the learner, loads the checkpoint, restores optimizer state when requested, and supports fp16 training. ------------------------------------------------------------------------ source ### resumed_pragma_dls ``` python 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.* ``` python 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 ``` python 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 ``` python 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 │ └─────────┴──────────┴────────┴─────────────────────┘ ``` python 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 │ └─────────┴─────┴────────┴────────────┴──────────┘ ``` python 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]

``` python n_keys, n_vals = tok.n_keys, tok.n_vals ``` ``` python 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, []) ``` python 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 ``` python 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') ``` ``` python 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) ``` ``` python import nbdev; nbdev.nbdev_export() ```