File size: 1,888 Bytes
4aeab24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def decrementer4(a3, a2, a1, a0, weights):
"""4-bit decrementer: returns (input - 1) mod 16"""
inp = torch.tensor([float(a3), float(a2), float(a1), float(a0)])
# Layer 1
y0 = int((inp @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
b1 = int((inp @ weights['b1.weight'].T + weights['b1.bias'] >= 0).item())
b2 = int((inp @ weights['b2.weight'].T + weights['b2.bias'] >= 0).item())
y1_and = int((inp @ weights['y1_and.weight'].T + weights['y1_and.bias'] >= 0).item())
# Layer 2
l2_in = torch.tensor([float(a3), float(a2), float(b1), float(b2), float(y1_and)])
y1 = int((l2_in @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
y2_or = int((l2_in @ weights['y2_or.weight'].T + weights['y2_or.bias'] >= 0).item())
y2_nand = int((l2_in @ weights['y2_nand.weight'].T + weights['y2_nand.bias'] >= 0).item())
y3_or = int((l2_in @ weights['y3_or.weight'].T + weights['y3_or.bias'] >= 0).item())
y3_nand = int((l2_in @ weights['y3_nand.weight'].T + weights['y3_nand.bias'] >= 0).item())
# Layer 3
l3_y2 = torch.tensor([float(y2_or), float(y2_nand)])
l3_y3 = torch.tensor([float(y3_or), float(y3_nand)])
y2 = int((l3_y2 @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
y3 = int((l3_y3 @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())
return [y3, y2, y1, y0]
if __name__ == '__main__':
w = load_model()
print('Decrementer4bit:')
for i in range(16):
a3, a2, a1, a0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
result = decrementer4(a3, a2, a1, a0, w)
out_val = result[0]*8 + result[1]*4 + result[2]*2 + result[3]
print(f' {i:2d} ({a3}{a2}{a1}{a0}) - 1 = {out_val:2d} ({"".join(map(str, result))})')
|