import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import gradio as gr

model = nn.Sequential(
            nn.Linear(11, 20),
            nn.ReLU(),
            nn.Linear(20, 5, bias=True))

PATH = "wine_model.pth"

model.load_state_dict(torch.load(PATH, weights_only=False))

def forward(model, input):
  preds = model(input)
  predicted_class = torch.argmax(preds, dim=-1) + 4
  return predicted_class

with gr.Blocks() as demo:
    gr.Markdown("Enter your wine data below:")
    input_df = gr.Dataframe(
        row_count=(2, "dynamic"),  # Allows adding/removing rows
        col_count=(11, "dynamic"),  # Allows adding/removing columns
        headers=list(df.columns)[:-1],
        label="Input Data",
        interactive=True,
        type="pandas" # Specify the desired input type for your function
    )
    output_text = gr.Textbox(label="Processed Output")

def process_data(input_dataframe):
    # Perform operations on the input_dataframe
    if isinstance(input_dataframe, pd.DataFrame):
        return forward(model, input_dataframe)
    return "Invalid input type"

input_df.change(fn=process_data, inputs=input_df, outputs=output_text)

demo.launch()