| import torch | |
| from safetensors.torch import load_file | |
| def load_model(path='model.safetensors'): | |
| return load_file(path) | |
| def ffs8(bits, w): | |
| """Find first set bit (1-indexed). Returns 0 if no bits set.""" | |
| inp = torch.tensor([float(b) for b in bits]) | |
| p = [] | |
| for i in range(8): | |
| pi = int((inp @ w[f'p{i}.weight'].T + w[f'p{i}.bias'] >= 0).item()) | |
| p.append(pi) | |
| f0 = 1 if (p[0] or p[2] or p[4] or p[6]) else 0 | |
| f1 = 1 if (p[1] or p[2] or p[5] or p[6]) else 0 | |
| f2 = 1 if (p[3] or p[4] or p[5] or p[6]) else 0 | |
| f3 = p[7] | |
| return f3, f2, f1, f0 | |
| if __name__ == '__main__': | |
| w = load_model() | |
| print('FFS8 selected tests:') | |
| print('Input | FFS | f3f2f1f0') | |
| print('---------+-----+---------') | |
| test_vals = [0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x06, 0xF0, 0xFF] | |
| for val in test_vals: | |
| bits = [(val >> j) & 1 for j in range(8)] | |
| f3, f2, f1, f0 = ffs8(bits, w) | |
| pos = 8*f3 + 4*f2 + 2*f1 + f0 | |
| print(f'{val:08b} | {pos} | {f3}{f2}{f1}{f0}') | |