import gradio as gr
import joblib
import pandas as pd
from huggingface_hub import hf_hub_download

# Download model from Hugging Face Hub
model_path = hf_hub_download(repo_id="abhishek/autotrain-iris-xgboost", filename="model.joblib")
model = joblib.load(model_path)

# Input labels expected by the model
feature_names = ['feat_SepalLengthCm', 'feat_SepalWidthCm', 'feat_PetalLengthCm', 'feat_PetalWidthCm']

def predict(sepal_length, sepal_width, petal_length, petal_width):
    data = pd.DataFrame([[sepal_length, sepal_width, petal_length, petal_width]], columns=feature_names)
    prediction = model.predict(data)[0]
    return f"Predicted Iris Class: {prediction}"

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Number(label="Sepal Length (cm)"),
        gr.Number(label="Sepal Width (cm)"),
        gr.Number(label="Petal Length (cm)"),
        gr.Number(label="Petal Width (cm)"),
    ],
    outputs=gr.Textbox(label="Prediction"),
    title="Iris Species Predictor 🌸",
    description="Enter flower features to predict the Iris species using a model trained with AutoTrain Tabular."
)

iface.launch()