model.py

import torch
from safetensors.torch import load_file

def load_model(path='model.safetensors'):
    return load_file(path)

def incrementer4(a3, a2, a1, a0, w):
    """Add 1 to 4-bit input (mod 16)."""
    inp = torch.tensor([float(a3), float(a2), float(a1), float(a0)])

    # Layer 1
    y0 = int((inp @ w['y0.weight'].T + w['y0.bias'] >= 0).item())
    c2 = int((inp @ w['c2.weight'].T + w['c2.bias'] >= 0).item())
    c3 = int((inp @ w['c3.weight'].T + w['c3.bias'] >= 0).item())
    y1_or = int((inp @ w['y1_or.weight'].T + w['y1_or.bias'] >= 0).item())
    y1_nand = int((inp @ w['y1_nand.weight'].T + w['y1_nand.bias'] >= 0).item())

    # Layer 2
    l2_in = torch.tensor([float(a3), float(a2), float(c2), float(c3), float(y1_or), float(y1_nand)])
    y1 = int((l2_in @ w['y1.weight'].T + w['y1.bias'] >= 0).item())
    y2_or = int((l2_in @ w['y2_or.weight'].T + w['y2_or.bias'] >= 0).item())
    y2_nand = int((l2_in @ w['y2_nand.weight'].T + w['y2_nand.bias'] >= 0).item())
    y3_or = int((l2_in @ w['y3_or.weight'].T + w['y3_or.bias'] >= 0).item())
    y3_nand = int((l2_in @ w['y3_nand.weight'].T + w['y3_nand.bias'] >= 0).item())

    # Layer 3
    y2 = int(y2_or + y2_nand - 2 >= 0)
    y3 = int(y3_or + y3_nand - 2 >= 0)

    return [y3, y2, y1, y0]

if __name__ == '__main__':
    w = load_model()
    print('incrementer4bit:')
    for i in range(16):
        a3, a2, a1, a0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
        result = incrementer4(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))})')