Update app.py

app.py

@@ -1,15 +1,16 @@
 import os, json, importlib.util, tempfile, traceback, torch, re, math
+import torch.nn as nn
 import torch.nn.functional as F
 import gradio as gr
 import pandas as pd
 import plotly.graph_objects as go
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModel
 
 # ===== ปรับได้จาก Settings > Variables & secrets ของ Space =====
 REPO_ID       = os.getenv("REPO_ID", "Dusit-P/thai-sentiment-wcb")
-DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "cnn_bilstm")  # หรือ "baseline"
+DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "cnn_bilstm")  # "cnn_bilstm" | "baseline" | "last4weighted_pure"
 HF_TOKEN      = os.getenv("HF_TOKEN", None)               # ถ้าโมเดลเป็น private ให้เพิ่ม secret ชื่อนี้
 
 # ---- theme colors (soft modern) ----
@@ -19,7 +20,7 @@ TEMPLATE  = "plotly_white"
 
 CACHE = {}
 
-# ---------- load architecture & weights from model repo ----------
+# ---------- โหลดสถาปัตยกรรมจาก repo (common/models.py) ----------
 def _import_models():
     if "models_module" in CACHE:
         return CACHE["models_module"]
@@ -30,20 +31,94 @@ def _import_models():
     CACHE["models_module"] = mod
     return mod
 
+# ---------- Fallback เผื่อ common/models.py ยังไม่รู้จัก Model3 ----------
+class _BaseHead(nn.Module):
+    def __init__(self, hidden_in, hidden_lstm=128, classes=2, dropout=0.3, pooling='masked_mean'):
+        super().__init__()
+        self.lstm = nn.LSTM(hidden_in, hidden_lstm, bidirectional=True, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(hidden_lstm*2, classes)
+        assert pooling in ['cls','masked_mean','masked_max']
+        self.pooling = pooling
+    def _pool(self, x, mask):
+        if self.pooling=='cls': return x[:,0,:]
+        mask = mask.unsqueeze(-1)
+        if self.pooling=='masked_mean':
+            s=(x*mask).sum(1); d=mask.sum(1).clamp(min=1e-6); return s/d
+        x=x.masked_fill(mask==0,-1e9); return x.max(1).values
+    def forward_after_bert(self, seq, mask):
+        x,_ = self.lstm(seq)
+        x = self._pool(x, mask)
+        return self.fc(self.dropout(x))
+
+class _Model3PureLast4(nn.Module):
+    """Last-4 weighted (Pure): LSTM รับ 768 จาก BERT"""
+    def __init__(self, base_model, hidden=128, classes=2, dropout=0.3, pooling='masked_mean'):
+        super().__init__()
+        self.bert = AutoModel.from_pretrained(base_model)
+        self.w = nn.Parameter(torch.ones(4))
+        H = self.bert.config.hidden_size
+        self.head = _BaseHead(H, hidden, classes, dropout, pooling)
+    def forward(self, ids, mask):
+        out = self.bert(input_ids=ids, attention_mask=mask, output_hidden_states=True)
+        last4 = out.hidden_states[-4:]
+        w = F.softmax(self.w, dim=0)
+        seq = sum(w[i]*last4[i] for i in range(4))  # [B,T,768]
+        return self.head.forward_after_bert(seq, mask)
+
+class _Model3ConvLast4(nn.Module):
+    """Last-4 weighted + Conv1d(→128): LSTM รับ 128"""
+    def __init__(self, base_model, hidden=128, classes=2, dropout=0.3, pooling='masked_mean'):
+        super().__init__()
+        self.bert = AutoModel.from_pretrained(base_model)
+        self.w = nn.Parameter(torch.ones(4))
+        H = self.bert.config.hidden_size
+        self.c1 = nn.Conv1d(H,128,3,padding=1)
+        self.c2 = nn.Conv1d(128,128,5,padding=2)
+        self.head = _BaseHead(128, hidden, classes, dropout, pooling)
+    def forward(self, ids, mask):
+        out = self.bert(input_ids=ids, attention_mask=mask, output_hidden_states=True)
+        last4 = out.hidden_states[-4:]
+        w = F.softmax(self.w, dim=0)
+        seq = sum(w[i]*last4[i] for i in range(4))      # [B,T,768]
+        x = F.relu(self.c1(seq.transpose(1,2)))
+        x = F.relu(self.c2(x)).transpose(1,2)           # [B,T,128]
+        return self.head.forward_after_bert(x, mask)
+
+def _create_model_fallback(arch: str, base_model: str):
+    """เลือกสถาปัตยกรรม fallback จากชื่อ arch ใน config.json"""
+    if arch in ("Model3_Pure_Last4Weighted", "last4weighted_pure", "last4_pure"):
+        return _Model3PureLast4(base_model)
+    if arch in ("Model3_MLP_Last4Weighted", "last4weighted"):
+        return _Model3ConvLast4(base_model)
+    raise ValueError(f"No fallback available for arch={arch}")
+
+# ---------- โหลดโมเดลจากโฟลเดอร์ใน repo (เช่น cnn_bilstm/, baseline/, last4weighted_pure/) ----------
 def load_model(model_name: str):
     key = f"model:{model_name}"
     if key in CACHE:
         return CACHE[key]
+
     cfg_path = hf_hub_download(REPO_ID, filename=f"{model_name}/config.json", token=HF_TOKEN)
     w_path   = hf_hub_download(REPO_ID, filename=f"{model_name}/model.safetensors", token=HF_TOKEN)
 
     with open(cfg_path, "r", encoding="utf-8") as f:
         cfg = json.load(f)
 
-    models = _import_models()
-    tok = AutoTokenizer.from_pretrained(cfg["base_model"])
-    model = models.create_model_by_name(cfg["arch"])
+    base_model = cfg.get("base_model", "airesearch/wangchanberta-base-att-spm-uncased")
+    arch_name  = cfg.get("arch", "")
+    tok = AutoTokenizer.from_pretrained(base_model)
+
+    # พยายามสร้างจาก common/models.py ก่อน ถ้าไม่สำเร็จค่อย fallback
+    try:
+        models = _import_models()
+        model = models.create_model_by_name(arch_name)
+    except Exception as e:
+        print(f"[INFO] Using fallback for arch={arch_name} ({e})")
+        model = _create_model_fallback(arch_name, base_model)
+
     state = load_file(w_path)
+    # ใช้ strict=True ถ้า key ตรง; ถ้าอยากกัน edge-case สามารถปรับเป็น strict=False ได้
     model.load_state_dict(state, strict=True)
     model.eval()
 
@@ -163,7 +238,7 @@ def _shop_summary(out_df: pd.DataFrame, max_shops=15):
     g = g.sort_values("total", ascending=False)
 
     table = g[["total","positive","negative"]].copy()
-    table["positive_rate(%)"] = (table["positive"] / table["total"] * 100).round(2)
+    table["positive_rate(%))"] = (table["positive"] / table["total"] * 100).round(2)
     table["negative_rate(%)"] = (table["negative"] / table["total"] * 100).round(2)
     table = table.reset_index().rename(columns={"index":"shop"})
 
@@ -211,7 +286,7 @@ def predict_one(text: str, model_choice: str):
         s = _norm_text(text)
         if not _is_substantive_text(s):
             return {"negative": 0.0, "positive": 0.0}, "invalid"
-        model_name = "baseline" if model_choice == "baseline" else "cnn_bilstm"
+        model_name = model_choice  # ใช้ชื่อโฟลเดอร์โดยตรง
         out = _predict_batch([s], model_name)[0]
         probs = {
             "negative": float(out["negative(%)"].rstrip("%"))/100.0,
@@ -225,7 +300,7 @@ def predict_one(text: str, model_choice: str):
 
 def predict_many(text_block: str, model_choice: str):
     try:
-        model_name = "baseline" if model_choice == "baseline" else "cnn_bilstm"
+        model_name = model_choice  # ใช้ชื่อโฟลเดอร์โดยตรง
         raw_lines = (text_block or "").splitlines()
         trimmed = [_norm_text(ln) for ln in raw_lines if _norm_text(ln)]
         cleaned, skipped = _clean_texts(trimmed)
@@ -257,7 +332,7 @@ def predict_csv(file_obj, model_choice: str, review_col_override: str = "", shop
         if file_obj is None:
             return pd.DataFrame(), None, go.Figure(), go.Figure(), go.Figure(), pd.DataFrame(), "กรุณาอัปโหลดไฟล์ CSV"
 
-        model_name = "baseline" if model_choice == "baseline" else "cnn_bilstm"
+        model_name = model_choice  # ใช้ชื่อโฟลเดอร์โดยตรง
         df = pd.read_csv(file_obj.name)
 
         auto_rev, auto_shop = _detect_cols(df)
@@ -333,10 +408,14 @@ def predict_csv(file_obj, model_choice: str, review_col_override: str = "", shop
         raise
 
 # ---------- Gradio UI ----------
+AVAILABLE_CHOICES = ["cnn_bilstm", "baseline", "last4weighted_pure"]  # เพิ่มชื่อโฟลเดอร์โมเดลใหม่ที่คุณอัปจริง
+if DEFAULT_MODEL not in AVAILABLE_CHOICES:
+    DEFAULT_MODEL = "cnn_bilstm"
+
 with gr.Blocks(title="Thai Sentiment API (Dusit-P)") as demo:
-    gr.Markdown("### Thai Sentiment (WangchanBERTa + LSTM/CNN Heads)")
+    gr.Markdown("### Thai Sentiment (WangchanBERTa + LSTM/CNN/Last4 Heads)")
 
-    model_radio = gr.Radio(choices=["cnn_bilstm","baseline"], value=DEFAULT_MODEL, label="เลือกโมเดล")
+    model_radio = gr.Radio(choices=AVAILABLE_CHOICES, value=DEFAULT_MODEL, label="เลือกโมเดล")
 
     with gr.Tab("Single"):
         t1 = gr.Textbox(lines=3, label="ข้อความรีวิว (1 ข้อความ)")