File size: 792 Bytes
dd97b24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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}')
|