Update predict_blend.py

predict_blend.py

@@ -1,21 +1,18 @@
 # predict_blend.py
 import os, json, numpy as np, pandas as pd, torch, lightgbm as lgb
 import torch.nn as nn
-from huggingface_hub import snapshot_download
-from pathlib import Path
-from itertools import product
-
-MODEL_REPO = "Antonio0616/FormingStar"
-
-# ✅ 반드시 dataset 으로 명시해야 함
 
-MODEL_DIR = snapshot_download(repo_id=MODEL_REPO, repo_type="dataset")
-ART_DIR = str(Path(MODEL_DIR).resolve())
+# =========================
+# Config
+# =========================
+from pathlib import Path
+BASE_DIR = Path(__file__).resolve().parent
+ART_DIR = str((BASE_DIR / "artifacts_blend").resolve())
 CAT_COL = "material"
 NUM_COLS = ["thickness","diameter","degree","upper_radius","lower_radius","LB","RB"]
 
 # =========================
-# FT-Transformer 정의
+# FT-Transformer
 # =========================
 class FTTransformer(nn.Module):
     def __init__(self, n_materials:int, n_num:int, d_model:int=192, nhead:int=8,
@@ -49,11 +46,35 @@ class FTTransformer(nn.Module):
         h = self.encoder(tokens)
         return self.head(h[:, 0, :])
 
+
 def _scale_like_fold(X_num: np.ndarray, mean: np.ndarray, scale: np.ndarray) -> np.ndarray:
     return ((X_num - mean) / scale).astype(np.float32)
 
 # =========================
-# 모델 불러오기 헬퍼
+# Material label helpers
+# =========================
+def _canonize_list(materials):
+    return [str(m).strip() for m in materials]
+
+def _build_alias2canon(canon_list):
+    alias2canon = {}
+    for c in canon_list:
+        alias2canon[c] = c
+        s = c.strip()
+        alias2canon[s] = c
+        if "." in s:
+            alias2canon[s.rstrip("0").rstrip(".")] = c
+        try:
+            v = float(s)
+            alias2canon[str(v)] = c
+            if v.is_integer():
+                alias2canon[str(int(v))] = c
+        except:
+            pass
+    return alias2canon
+
+# =========================
+# Loader helpers
 # =========================
 def _first_existing(*paths):
     for p in paths:
@@ -108,6 +129,15 @@ def _load_json_like(art_dir: str, basename: str) -> dict:
     with open(p, "r", encoding="utf-8") as f:
         return json.load(f)
 
+def _load_materials(art_dir: str, folds_ft):
+    try:
+        return _load_json_like(art_dir, "materials")["materials"]
+    except FileNotFoundError:
+        return folds_ft[0]["materials"]
+
+def _load_best_alpha(art_dir: str) -> float:
+    return float(_load_json_like(art_dir, "blend_alpha")["best_alpha"])
+
 # =========================
 # Predictor
 # =========================
@@ -116,22 +146,40 @@ class BlendPredictor:
         self.art_dir = art_dir
         self.folds_ft = _load_ft_folds(art_dir)
         self.boosters = _load_lgbm_folds(art_dir)
-        self.materials = _load_json_like(art_dir, "materials")["materials"]
-        self.best_alpha = float(_load_json_like(art_dir, "blend_alpha")["best_alpha"])
+        self.materials = _load_materials(art_dir, self.folds_ft)
+        self.best_alpha = _load_best_alpha(art_dir)
 
-        self.materials = [str(m).strip() for m in self.materials]
-        self.mat2id = {m: i for i, m in enumerate(self.materials)}
+        self.materials_canon = _canonize_list(self.materials)
+        self.alias2canon = _build_alias2canon(self.materials_canon)
+        self.mat2id = {m: i for i, m in enumerate(self.materials_canon)}
         self.unknown_policy = unknown_policy
 
     def _prep_df(self, df_new: pd.DataFrame) -> pd.DataFrame:
         df = df_new.copy()
+        need = [CAT_COL] + NUM_COLS
+        missing = [c for c in need if c not in df.columns]
+        if missing:
+            raise ValueError(f"Missing columns in input: {missing}")
+
         df[CAT_COL] = df[CAT_COL].astype(str).str.strip()
+        df["_mat_canon"] = df[CAT_COL].map(self.alias2canon)
+
         if self.unknown_policy == "error":
-            unknown = df.loc[~df[CAT_COL].isin(self.materials), CAT_COL].unique().tolist()
+            unknown = df.loc[df["_mat_canon"].isna(), CAT_COL].unique().tolist()
             if unknown:
-                raise ValueError(f"Unknown materials in input {unknown}")
-        df["_mat_id"] = df[CAT_COL].map(self.mat2id).fillna(0).astype(int)
+                raise ValueError(
+                    f"Unknown materials in input {unknown}. "
+                    f"Known materials: {self.materials_canon[:10]}{' ...' if len(self.materials_canon)>10 else ''}"
+                )
+            df["_mat_id"] = df["_mat_canon"].map(self.mat2id).astype(int)
+        else:
+            df["_mat_canon"] = df["_mat_canon"].fillna(self.materials_canon[0])
+            df["_mat_id"] = df["_mat_canon"].map(self.mat2id).astype(int)
+
         df[NUM_COLS] = df[NUM_COLS].apply(pd.to_numeric, errors="coerce")
+        if df[NUM_COLS].isnull().any().any():
+            bad = df[NUM_COLS].columns[df[NUM_COLS].isnull().any()].tolist()
+            raise ValueError(f"Non-numeric values detected in columns: {bad}")
         return df
 
     def predict_ft(self, df_new: pd.DataFrame) -> np.ndarray:
@@ -150,7 +198,7 @@ class BlendPredictor:
     def predict_lgbm(self, df_new: pd.DataFrame) -> np.ndarray:
         df = self._prep_df(df_new)
         X = df[[CAT_COL] + NUM_COLS].copy()
-        X[CAT_COL] = pd.Categorical(df[CAT_COL], categories=self.materials)
+        X[CAT_COL] = pd.Categorical(df["_mat_canon"], categories=self.materials_canon)
         preds = [bst.predict(X, num_iteration=getattr(bst, "best_iteration", None))
                  for bst in self.boosters]
         return np.mean(preds, axis=0)
@@ -162,32 +210,24 @@ class BlendPredictor:
         p_lgb = self.predict_lgbm(df_new)
         return alpha * p_dl + (1 - alpha) * p_lgb
 
-    # ✅ 범위 입력 → 경우의 수 생성 지원
-    def expand_ranges(self, cfg: dict) -> pd.DataFrame:
-        # cfg: {"material": ["SPCC"], "min_thickness": 0.7, "max_thickness": 1.2, "thickness_step": 0.1, ...}
-        keys = []
-        values = []
-
-        # 범주형
-        keys.append("material")
-        values.append(cfg["materials"])
-
-        # 연속형
-        for col in ["thickness","diameter","degree","upper_radius","lower_radius"]:
-            lo = cfg[f"min_{col}"]
-            hi = cfg[f"max_{col}"]
-            step = cfg[f"{col}_step"]
-            values.append(np.arange(lo, hi+1e-9, step).round(3))
-            keys.append(col)
-
-        # 비드 (LB, RB 변환)
-        bead_map = {
-            "none": (0,0), "left": (1,0), "right": (0,1), "double": (1,1)
-        }
-        beads = [bead_map[b] for b in cfg.get("beads", ["none"])]
-        LB, RB = zip(*beads)
-        keys.extend(["LB","RB"])
-        values.extend([LB, RB])
-
-        combos = [dict(zip(keys, v)) for v in product(*values)]
-        return pd.DataFrame(combos)
+# =========================
+# Example run
+# =========================
+if __name__ == "__main__":
+    base = {
+        "thickness": 1, "diameter": 20, "degree": 73,
+        "upper_radius": 3, "lower_radius": 2,
+        "LB": 0, "RB": 1,
+    }
+    df_new = pd.DataFrame([
+        {**base, "material": "590"},
+        {**base, "material": "440"},
+    ])
+
+    predictor = BlendPredictor(ART_DIR, unknown_policy="error")
+    print("materials (trained):", predictor.materials_canon[:10])
+    print("best_alpha:", predictor.best_alpha)
+
+    print("\nDL only :", predictor.predict_blend(df_new, alpha=1.0))
+    print("LGBM only:", predictor.predict_blend(df_new, alpha=0.0))
+    print("Blend    :", predictor.predict_blend(df_new))