import torch
from safetensors.torch import load_file

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

def half_subtractor(a, b, weights):
    """Half subtractor: computes a - b, returns (diff, borrow)"""
    inp = torch.tensor([float(a), float(b)])
    # XOR for diff
    l1 = (inp @ weights['xor.layer1.weight'].T + weights['xor.layer1.bias'] >= 0).float()
    diff = int((l1 @ weights['xor.layer2.weight'].T + weights['xor.layer2.bias'] >= 0).item())
    # Borrow
    borrow = int((inp @ weights['borrow.weight'].T + weights['borrow.bias'] >= 0).item())
    return diff, borrow

if __name__ == '__main__':
    w = load_model()
    print('HalfSubtractor truth table:')
    for a in [0, 1]:
        for b in [0, 1]:
            diff, borrow = half_subtractor(a, b, w)
            print(f'{a} - {b} = {diff}, borrow={borrow}')