File size: 2,595 Bytes
088ff92 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def hamming_distance(a3, a2, a1, a0, b3, b2, b1, b0, weights):
"""Compute Hamming distance between two 4-bit values. Returns (d2, d1, d0)."""
inp = torch.tensor([float(a3), float(a2), float(a1), float(a0),
float(b3), float(b2), float(b1), float(b0)])
# Layer 1: diff_hi and diff_lo for each bit
diff_hi, diff_lo = [], []
for i in range(4):
dh = int((inp @ weights[f'layer1.diff_hi_{i}.weight'].T + weights[f'layer1.diff_hi_{i}.bias'] >= 0).item())
dl = int((inp @ weights[f'layer1.diff_lo_{i}.weight'].T + weights[f'layer1.diff_lo_{i}.bias'] >= 0).item())
diff_hi.append(dh)
diff_lo.append(dl)
# Layer 2: XOR = OR(diff_hi, diff_lo)
diffs = []
for i in range(4):
d_inp = torch.tensor([float(diff_hi[i]), float(diff_lo[i])])
d = int((d_inp @ weights[f'layer2.diff_{i}.weight'].T + weights[f'layer2.diff_{i}.bias'] >= 0).item())
diffs.append(d)
# Layer 3: threshold detectors
diff_tensor = torch.tensor([float(d) for d in diffs])
ge = []
for k in range(1, 5):
g = int((diff_tensor @ weights[f'layer3.ge{k}.weight'].T + weights[f'layer3.ge{k}.bias'] >= 0).item())
ge.append(g)
# Layer 4: binary encoding
ge_tensor = torch.tensor([float(g) for g in ge])
d2 = int((ge_tensor @ weights['layer4.d2.weight'].T + weights['layer4.d2.bias'] >= 0).item())
d1 = int((ge_tensor @ weights['layer4.d1.weight'].T + weights['layer4.d1.bias'] >= 0).item())
d0_part1 = int((ge_tensor @ weights['layer4.d0_part1.weight'].T + weights['layer4.d0_part1.bias'] >= 0).item())
d0_part2 = int((ge_tensor @ weights['layer4.d0_part2.weight'].T + weights['layer4.d0_part2.bias'] >= 0).item())
# Layer 5: d0 = OR(d0_part1, d0_part2)
d0_inp = torch.tensor([float(d0_part1), float(d0_part2)])
d0 = int((d0_inp @ weights['layer5.d0.weight'].T + weights['layer5.d0.bias'] >= 0).item())
return d2, d1, d0
if __name__ == '__main__':
w = load_model()
print('Hamming Distance 4-bit examples:')
test_cases = [(0b0000, 0b0000), (0b1111, 0b0000), (0b1010, 0b0101), (0b1100, 0b1010)]
for a, b in test_cases:
a3, a2, a1, a0 = (a >> 3) & 1, (a >> 2) & 1, (a >> 1) & 1, a & 1
b3, b2, b1, b0 = (b >> 3) & 1, (b >> 2) & 1, (b >> 1) & 1, b & 1
d2, d1, d0 = hamming_distance(a3, a2, a1, a0, b3, b2, b1, b0, w)
dist = 4*d2 + 2*d1 + d0
print(f' HD({a:04b}, {b:04b}) = {dist}')
|