import gradio as gr
import torch
import torch.nn as nn


class XORModel(nn.Module):
    def __init__(self):
        super(XORModel, self).__init__()
        self.fc1 = nn.Linear(2, 2)
        self.fc2 = nn.Linear(2, 1)

    def forward(self, x):
        x = torch.sigmoid(self.fc1(x))
        x = torch.sigmoid(self.fc2(x))
        return x


model = XORModel()
model.load_state_dict(torch.load("modelo_xor.pt"))
model.eval()


def predict(x1, x2):
    resultado = 0
    inputs = torch.tensor([[x1, x2]], dtype=torch.float32)
    with torch.no_grad():
        prediction = model(inputs)
        prediction = (prediction > 0.5).item()
        if prediction == 0:
            resultado = 0
        else:
            resultado = 1
    return resultado


# Define la interfaz Gradio
iface = gr.Interface(
    fn=predict,
    inputs=[gr.components.Number(), gr.components.Number()],
    outputs=gr.components.Label(),
    title="Modelo XOR con PyTorch",
    description="Ingresa dos números (0 o 1) para predecir el resultado XOR.",
    live=False,  # Muestra el resultado en tiempo real
)

iface.launch()