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