Upload app.py with huggingface_hub

app.py

@@ -1,11 +1,13 @@
-
-# !pip install -q autogluon.tabular gradio huggingface_hub pandas
+import pathlib, shutil, zipfile, os, traceback
+import pandas as pd
+import gradio as gr
 
 from huggingface_hub import hf_hub_download
 from autogluon.tabular import TabularPredictor
-import pathlib, shutil, zipfile
-import pandas as pd
-import gradio as gr
+
+# ---------------- UI copy ----------------
+TITLE = "🧱 LEGO Brick Classifier"
+DESC  = "Predicts whether a LEGO piece is Standard, Flat, or Sloped from basic dimensions."
 
 # ---------------- Settings ----------------
 MODEL_REPO_ID = "Iris314/classical-automl-model"
@@ -28,9 +30,12 @@ def load_predictor():
     local_zip = hf_hub_download(
         repo_id=MODEL_REPO_ID,
         filename=ZIP_FILENAME,
-        repo_type="model"
+        repo_type="model",
+        local_dir=str(CACHE_DIR),
+        local_dir_use_symlinks=False,
     )
-    if EXTRACT_DIR.exists(): shutil.rmtree(EXTRACT_DIR)
+    if EXTRACT_DIR.exists():
+        shutil.rmtree(EXTRACT_DIR)
     EXTRACT_DIR.mkdir(parents=True)
     with zipfile.ZipFile(local_zip, "r") as zf:
         zf.extractall(EXTRACT_DIR)
@@ -38,7 +43,11 @@ def load_predictor():
     path = kids[0] if len(kids) == 1 and kids[0].is_dir() else EXTRACT_DIR
     return TabularPredictor.load(str(path), require_py_version_match=False)
 
-PREDICTOR = load_predictor()
+try:
+    PREDICTOR = load_predictor()
+except Exception as e:
+    PREDICTOR = None
+    print("Failed to load predictor:", e)
 
 # ---------------- Helpers ----------------
 def _cast_and_rename(row_dict):
@@ -46,20 +55,21 @@ def _cast_and_rename(row_dict):
     row["Length"] = float(row["Length"])
     row["Height"] = float(row["Height"])
     row["Width"]  = float(row["Width"])
-    row["Studs"]  = int(round(float(row["Studs"])))  # gr.Number returns float
+    # gr.Number returns float; round & cast for integer feature
+    row["Studs"]  = int(round(float(row["Studs"])))
     X_ui = pd.DataFrame([row], columns=FEATURE_COLS_UI)
     X_model = X_ui.rename(columns=COLUMN_ALIAS)
     return X_model
 
 def classify_brick(length, height, width, studs):
     try:
+        if PREDICTOR is None:
+            raise RuntimeError("Model failed to load on startup. Check model artifact path & runtime deps.")
+
         X = _cast_and_rename({
             "Length": length, "Height": height, "Width": width, "Studs": studs
         })
 
-        pred = PREDICTOR.predict(X)
-        pred_val = pred.iloc[0] if hasattr(pred, "iloc") else pred
-
         # Try probabilities; fall back to label
         try:
             proba = PREDICTOR.predict_proba(X)
@@ -68,29 +78,30 @@ def classify_brick(length, height, width, studs):
             s.index = [str(k) for k in s.index]  # ensure JSON-serializable keys
             return {k: float(v) for k, v in s.items()}
         except Exception:
+            pred = PREDICTOR.predict(X)
+            pred_val = pred.iloc[0] if hasattr(pred, "iloc") else pred
             return {"prediction": str(pred_val)}
     except Exception as e:
-        import traceback
-        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()}
-
-# ---------------- Quick test (uses correct names) ----------------
-# test_X = _cast_and_rename({"Length": 4, "Height": 1.2, "Width": 2, "Studs": 4})
-# print("Prediction:", PREDICTOR.predict(test_X))
-# print("Probabilities:\n", PREDICTOR.predict_proba(test_X))
+        return {
+            "error": f"{type(e).__name__}: {e}",
+            "traceback": traceback.format_exc(limit=1)
+        }
 
-# ---------------- Gradio (Interface version) ----------------
+# ---------------- Gradio ----------------
 demo = gr.Interface(
     fn=classify_brick,
     inputs=[
-        gr.Slider(1, 10, step=0.1, value=4,  label="Length"),
+        gr.Slider(1, 10, step=0.1, value=4,   label="Length"),
         gr.Slider(0.2, 5, step=0.1, value=1.2, label="Height"),
-        gr.Slider(1, 10, step=0.1, value=2,  label="Width"),
+        gr.Slider(1, 10, step=0.1, value=2,   label="Width"),
         gr.Number(value=4, precision=0, label="Studs"),
     ],
     outputs=gr.Label(num_top_classes=3, label="Predicted Class / Probabilities"),
     examples=[[4, 1.2, 2, 4], [2, 0.6, 2, 2], [3, 2.0, 2, 2]],
-    title="🧱 LEGO Brick Classifier",
-    description="Predicts whether a LEGO piece is Standard, Flat, or Sloped."
+    title=TITLE,
+    description=DESC
 )
 
-demo.launch()
+if __name__ == "__main__":
+    # In Spaces, no share=True needed
+    demo.launch()