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

# Define the model architecture
class SumModel(nn.Module):
    def __init__(self):
        super(SumModel, self).__init__()
        self.fc1 = nn.Linear(2, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 1)

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

# Load the pre-trained model
model = SumModel()
model.load_state_dict(torch.load('sum_model.pth'))
model.eval()  # Set the model to evaluation mode

# Function to predict the sum of two numbers
def calculate_sum(num1, num2):
    # Prepare the input tensor
    inputs = torch.tensor([[num1, num2]]).float()
    
    # Forward pass
    with torch.no_grad():  # Disable gradient tracking during inference
        outputs = model(inputs)
    
    # Get the predicted sum
    predicted_sum = outputs.item()
    
    return predicted_sum

# Create a Gradio interface
iface = gr.Interface(
    fn=calculate_sum,
    inputs=["number", "number"],
    outputs="number",
    title="Sum Predictor",
    description="Enter two numbers to predict their sum"
)

# Launch the Gradio app
iface.launch()