"""Compare original full BF16 projection with the extracted row projection."""
import json
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import load_file

root = Path('/model')
index = json.loads((root / 'model.safetensors.index.json').read_text())['weight_map']
with safe_open(root / index['lm_head.weight'], framework='pt', device='cpu') as file:
    full = file.get_tensor('lm_head.weight').to('cuda')
head = load_file('/head/head.safetensors', device='cuda')
ids = head['token_ids']
weight = head['weight']
torch.testing.assert_close(full.index_select(0, ids), weight, rtol=0, atol=0)
torch.manual_seed(0)
checks = []
for batch in (1, 4, 16):
    hidden = torch.randn(batch, full.shape[1], device='cuda', dtype=torch.bfloat16)
    reference = torch.nn.functional.linear(hidden, full).index_select(-1, ids)
    actual = torch.nn.functional.linear(hidden, weight)
    torch.testing.assert_close(actual, reference, rtol=0.01, atol=0.125)
    checks.append({'batch': batch,
                   'max_abs_difference': (actual.float() - reference.float()).abs().max().item(),
                   'argmax_matches': bool((actual.argmax(-1) == reference.argmax(-1)).all())})
print(json.dumps({'weights_exact': True, 'full_shape': list(full.shape),
                  'linear_shape': list(weight.shape), 'checks': checks}, indent=2))