IrisModelProd / app.py
DavidMoreda's picture
Subida app
c1e3ba9
# Import the necessary libraries
import gradio as gr
import joblib
import numpy as np
from huggingface_hub import hf_hub_download
# Download both models from HF Hub at startup
# Both models live in the same repository: brjapon/iris-dt
dt_path = hf_hub_download(repo_id="DavidMoreda/iris-dt", filename="iris_dt.joblib", repo_type="model")
lr_path = hf_hub_download(repo_id="DavidMoreda/iris-dt", filename="iris_logreg.joblib", repo_type="model")
models = {
"Decision Tree": joblib.load(dt_path),
"Logistic Regression": joblib.load(lr_path),
}
LABELS = {0: "Iris-setosa", 1: "Iris-versicolor", 2: "Iris-virginica"}
# The function now accepts a model_choice parameter (from a Gradio dropdown)
def predict_iris(model_choice, sepal_length, sepal_width, petal_length, petal_width):
pipeline = models[model_choice]
input = np.array([[sepal_length, sepal_width, petal_length, petal_width]])
prediction = pipeline.predict(input)
return LABELS.get(int(prediction[0]), "Invalid prediction")
interface = gr.Interface(
fn=predict_iris,
inputs=[
gr.Dropdown(choices=["Decision Tree", "Logistic Regression"], label="Model", value="Decision Tree"),
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="text",
live=True,
title="Iris Species Identifier",
description="Choose a model and enter the four measurements to predict the Iris species.",
flagging_mode="manual",
flagging_dir="flagged"
)
if __name__ == "__main__":
interface.launch()
'''
# The Flag button allows users (or testers) to mark or “flag”
# a particular input-output interaction for later review.
# When someone clicks Flag, Gradio saves the input values (and often the output) to a log.csv file
# letting you keep track of interesting or potentially problematic cases for debugging or analysis later on
'''