import gradio as gr
from model import SimpleNN, predict
import torch
from sklearn.preprocessing import LabelEncoder


model = SimpleNN(input_size=4, hidden_size=64, output_size=5)
model.load_state_dict(torch.load("your_model.pth"))
model.eval()


label_encoder = LabelEncoder()


def classifier(rarity, size, color, country):
    label_encoder.fit([rarity, color, country])

    input_data = [
        label_encoder.transform([rarity])[0],
        size,
        label_encoder.transform([color])[0],
        label_encoder.transform([country])[0],
    ]

    predicted_class_index = predict(model, input_data)

    predicted_species = label_encoder.inverse_transform([predicted_class_index])[0]

    return (
        f"{rarity} {size} {color} {country} => Predicted Species: {predicted_species}"
    )


iface = gr.Interface(
    fn=classifier, inputs=["text", "number", "text", "text"], outputs="text"
)


iface.launch(share=True)