"""Offline load/forward/generation/round-trip checks of the exported checkpoint."""
import hashlib
import json
from pathlib import Path
import torch
from safetensors import safe_open
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

root = Path(__file__).resolve().parent
path = root/'model'
torch.set_num_threads(2)
tok = AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False)
model = AutoModelForCausalLM.from_pretrained(path, local_files_only=True,
    trust_remote_code=False, torch_dtype=torch.float32).eval()
count = sum(p.numel() for p in model.parameters())
assert count == 49999872
assert len(tok) == 23258
with safe_open(str(path/'model.safetensors'), framework='pt') as f:
    tensor_count = len(f.keys())
    assert all(f.get_tensor(k).dtype == torch.float16 for k in f.keys())
    assert all(torch.isfinite(f.get_tensor(k)).all() for k in f.keys())
samples = []
for text in ['Однажды вечером', 'Once upon a time']:
    assert tok.decode(tok.encode(text)) == text
    set_seed(42)
    inputs = tok(text, return_tensors='pt')
    with torch.inference_mode():
        logits = model(**inputs).logits
        assert torch.isfinite(logits).all()
        generated = model.generate(**inputs, max_new_tokens=48, do_sample=True,
            temperature=.8, top_p=.9, repetition_penalty=1.1,
            pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
    samples.append({'prompt':text,'output':tok.decode(generated[0],skip_special_tokens=True)})
report = {'offline_transformers_load':'passed','trust_remote_code':False,
    'parameters':count,'vocabulary_size':len(tok),'tensor_count':tensor_count,
    'weights_dtype':'float16','finite_weights_and_logits':True,
    'russian_english_tokenizer_roundtrip':'passed',
    'generation_test':'execution only, NOT linguistic-quality certification',
    'generation_seed':42,'max_new_tokens':48,'temperature':.8,'top_p':.9,
    'repetition_penalty':1.1,'uncurated_samples':samples,
    'safetensors_bytes':(path/'model.safetensors').stat().st_size,
    'safetensors_sha256':hashlib.sha256((path/'model.safetensors').read_bytes()).hexdigest()}
(path/'verification.json').write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf-8')
print(json.dumps(report,ensure_ascii=False,indent=2))