import sys
sys.path.insert(0, '.')
import torch
from prepare import Tokenizer, evaluate_bpb
from train import detect_runtime, build_model_config, GPT
runtime = detect_runtime()
tokenizer = Tokenizer.from_directory()
vocab_size = tokenizer.get_vocab_size()
regime = type('Regime', (), {'depth':8,'aspect_ratio':64,'window_pattern':'SSSL'})()
config = build_model_config(regime, vocab_size, runtime, use_activation_checkpointing=False)
with torch.device('meta'):
    model = GPT(config)
model.to_empty(device=runtime.device)
state = torch.load('checkpoint_pre_eval.pt', map_location=runtime.device)
model.load_state_dict(state)
model.eval()
with torch.amp.autocast(device_type=runtime.device_type, dtype=runtime.amp_dtype):
    val_bpb = evaluate_bpb(model, tokenizer, batch_size=4, device=runtime.device, dataset=tokenizer.dataset, eval_tokens=65536)
print(f'val_bpb_small_eval={val_bpb:.6f}')