Update predict_blend.py

predict_blend.py

@@ -1,21 +1,23 @@
 # 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 huggingface_hub import snapshot_download   # ✅ Hugging Face dataset 불러오기
+from itertools import product
 
 # =========================
 # Config
 # =========================
-# ✅ Hugging Face Dataset에서 모델 파일 자동 다운로드
-dataset_path = snapshot_download(repo_id="Antonio0616/foemingstar-model")
-ART_DIR = os.path.join(dataset_path, "")   # artifacts_blend 폴더 대신 Dataset 사용
+MODEL_REPO = "Antonio0616/foemingstar-model"   # 제자님 repo ID
+BASE_DIR = Path(__file__).resolve().parent
+MODEL_DIR = snapshot_download(repo_id=MODEL_REPO)  # Hugging Face에서 모델 다운로드
+ART_DIR = str(Path(MODEL_DIR).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,10 +51,11 @@ class FTTransformer(nn.Module):
         h = self.encoder(tokens)
         return self.head(h[:, 0, :])
 
-# ✅ 이하 부분은 그대로 유지 (Loader, BlendPredictor 등)
+def _scale_like_fold(X_num: np.ndarray, mean: np.ndarray, scale: np.ndarray) -> np.ndarray:
+    return ((X_num - mean) / scale).astype(np.float32)
 
 # =========================
-# Loader helpers
+# 모델 불러오기 헬퍼
 # =========================
 def _first_existing(*paths):
     for p in paths:
@@ -107,15 +110,6 @@ 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
 # =========================
@@ -124,40 +118,22 @@ 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_materials(art_dir, self.folds_ft)
-        self.best_alpha = _load_best_alpha(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_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.materials = [str(m).strip() for m in self.materials]
+        self.mat2id = {m: i for i, m in enumerate(self.materials)}
         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["_mat_canon"].isna(), CAT_COL].unique().tolist()
+            unknown = df.loc[~df[CAT_COL].isin(self.materials), CAT_COL].unique().tolist()
             if unknown:
-                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)
-
+                raise ValueError(f"Unknown materials in input {unknown}")
+        df["_mat_id"] = df[CAT_COL].map(self.mat2id).fillna(0).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:
@@ -176,7 +152,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["_mat_canon"], categories=self.materials_canon)
+        X[CAT_COL] = pd.Categorical(df[CAT_COL], categories=self.materials)
         preds = [bst.predict(X, num_iteration=getattr(bst, "best_iteration", None))
                  for bst in self.boosters]
         return np.mean(preds, axis=0)
@@ -188,24 +164,32 @@ class BlendPredictor:
         p_lgb = self.predict_lgbm(df_new)
         return alpha * p_dl + (1 - alpha) * p_lgb
 
-# =========================
-# 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))
+    # ✅ 범위 입력 → 경우의 수 생성 지원
+    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)