diff options
Diffstat (limited to 'ep_run/casc_gen.py')
| -rw-r--r-- | ep_run/casc_gen.py | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/ep_run/casc_gen.py b/ep_run/casc_gen.py index e535df8..1d88f7c 100644 --- a/ep_run/casc_gen.py +++ b/ep_run/casc_gen.py @@ -12,7 +12,10 @@ ap.add_argument('--topk', type=int, default=40) ap.add_argument('--prompt', default='Once upon a time') args = ap.parse_args() -DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe') +import sys as _sys, torch as _t +_ckpath = next((a for a in _sys.argv if a.endswith('.pt')), 'runs/stage1b_ep_muon_s55000.pt') +_cfg = _t.load(_ckpath, map_location='cpu', weights_only=False).get('config', {}) +DD = Path('/home/yurenh2/ept/ep_run/data') / _cfg.get('data', 'tinystories_bpe') vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size'] from tokenizers import Tokenizer tk = Tokenizer.from_file(str(DD / 'tokenizer.json')) |