File size: 1,702 Bytes
dc1dc20 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from __future__ import annotations
from pathlib import Path
import gradio as gr
import joblib
import pandas as pd
ARTIFACT_PATH = Path(__file__).with_name("model.joblib")
def _load_bundle() -> dict:
if not ARTIFACT_PATH.exists():
from train import train_and_save
train_and_save(ARTIFACT_PATH)
return joblib.load(ARTIFACT_PATH)
BUNDLE = _load_bundle()
MODEL = BUNDLE["model"]
TARGET_NAMES = BUNDLE["target_names"]
FEATURE_NAMES = BUNDLE["feature_names"]
def predict(sepal_length: float, sepal_width: float, petal_length: float, petal_width: float):
x = pd.DataFrame(
[[sepal_length, sepal_width, petal_length, petal_width]],
columns=FEATURE_NAMES,
)
pred_idx = int(MODEL.predict(x)[0])
pred_label = TARGET_NAMES[pred_idx]
if hasattr(MODEL, "predict_proba"):
proba = MODEL.predict_proba(x)[0]
proba_dict = {str(TARGET_NAMES[i]): float(proba[i]) for i in range(len(TARGET_NAMES))}
else:
proba_dict = {TARGET_NAMES[pred_idx]: 1.0}
return pred_label, proba_dict
demo = gr.Interface(
fn=predict,
inputs=[
gr.Number(label=FEATURE_NAMES[0], value=5.8),
gr.Number(label=FEATURE_NAMES[1], value=3.0),
gr.Number(label=FEATURE_NAMES[2], value=4.0),
gr.Number(label=FEATURE_NAMES[3], value=1.2),
],
outputs=[
gr.Textbox(label="Predicted class"),
gr.Label(label="Class probabilities"),
],
title="Iris Classification (KNN)",
description=(
"KNN classifier trained on the classic Iris dataset. "
"Enter measurements and get a predicted species + probabilities."
),
)
if __name__ == "__main__":
demo.launch()
|