import torch
from safetensors.torch import load_file

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

def mux(a, b, s, weights):
    """2:1 Multiplexer: returns a if s=0, b if s=1"""
    inp = torch.tensor([float(a), float(b), float(s)])
    l1 = (inp @ weights['layer1.weight'].T + weights['layer1.bias'] >= 0).float()
    out = (l1 @ weights['layer2.weight'].T + weights['layer2.bias'] >= 0).float()
    return int(out.item())

if __name__ == '__main__':
    w = load_model()
    print('MUX truth table:')
    for a in [0, 1]:
        for b in [0, 1]:
            for s in [0, 1]:
                result = mux(a, b, s, w)
                print(f'MUX({a}, {b}, s={s}) = {result}')