import joblib
import os

import gradio as gr

from huggingface_hub import hf_hub_download


file_path = hf_hub_download("osanseviero/wine-quality", "sklearn_model.joblib",
                            use_auth_token=os.environ['TOKEN'])
model = joblib.load(file_path)

def predict(data):
    return model.predict(data.to_numpy())


headers = [
    "fixed acidity",
    "volatile acidity",
    "citric acid",
    "residual sugar",
    "chlorides",
    "free sulfur dioxide",
    "total sulfur dioxide",
    "density",
    "pH",
    "sulphates",
    "alcohol",
]
default = [
    [7.4, 0.7, 0, 1.9, 0.076, 11, 34, 0.9978, 3.51, 0.56, 9.4],
    [7.8, 0.88, 0, 2.6, 0.098, 25, 67, 0.9968, 3.2, 0.68, 9.8],
    [7.8, 0.76, 0.04, 2.3, 0.092, 15, 54, 0.997, 3.26, 0.65, 9.8],
]

iface = gr.Interface(
    predict,
    title="Wine Quality predictor with SKLearn",
    inputs=gr.inputs.Dataframe(
        headers=headers,
        default=default,
    ),
    outputs="numpy",
    description="Learn how to create demos of private models at https://huggingface.co/spaces/osanseviero/tips-and-tricks"
)
iface.launch()