| 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}') | |