import torch
from safetensors.torch import load_file

def load_model(path='model.safetensors'):
    return load_file(path)

def onehot_decode(y3, y2, y1, y0, weights):
    """Convert 4-bit one-hot encoding to 2-bit binary."""
    inp = torch.tensor([float(y3), float(y2), float(y1), float(y0)])

    a0 = int((inp @ weights['a0.weight'].T + weights['a0.bias'] >= 0).item())
    a1 = int((inp @ weights['a1.weight'].T + weights['a1.bias'] >= 0).item())

    return a1, a0

if __name__ == '__main__':
    w = load_model()
    print('One-Hot Decoder (4-to-2):')
    test_cases = [(0,0,0,1), (0,0,1,0), (0,1,0,0), (1,0,0,0)]
    for y3, y2, y1, y0 in test_cases:
        a1, a0 = onehot_decode(y3, y2, y1, y0, w)
        val = a1 * 2 + a0
        print(f'  {y3}{y2}{y1}{y0} -> {a1}{a0} = {val}')