| 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() |
|
|
