Update common/models.py

common/models.py

@@ -1,58 +1,125 @@
-# common/models.py
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers import AutoModel
-
-# ตั้งค่าพื้นฐานให้ตรงกับตอนเทรน
-BASE_MODEL_NAME = "airesearch/wangchanberta-base-att-spm-uncased"
-POOLING_AFTER_LSTM = "masked_mean"
-
-class BaseHead(nn.Module):
-    def __init__(self, hidden_in, hidden_lstm=128, num_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, num_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 Model1Baseline(nn.Module):
-    def __init__(self, name=BASE_MODEL_NAME, hidden=128, dropout=0.3, classes=2, pooling=POOLING_AFTER_LSTM):
-        super().__init__()
-        self.bert = AutoModel.from_pretrained(name)
-        self.head = BaseHead(self.bert.config.hidden_size, hidden, classes, dropout, pooling)
-    def forward(self, ids, mask):
-        out = self.bert(input_ids=ids, attention_mask=mask)
-        return self.head.forward_after_bert(out.last_hidden_state, mask)
-
-class Model2CNNBiLSTM(nn.Module):
-    def __init__(self, name=BASE_MODEL_NAME, hidden=128, dropout=0.3, classes=2, pooling=POOLING_AFTER_LSTM):
-        super().__init__()
-        self.bert = AutoModel.from_pretrained(name)
-        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).last_hidden_state
-        x = F.relu(self.c1(out.transpose(1,2)))
-        x = F.relu(self.c2(x)).transpose(1,2)
-        return self.head.forward_after_bert(x, mask)
-
-def create_model_by_name(model_name):
-    if model_name == "Model1_Baseline":
-        return Model1Baseline()
-    elif model_name == "Model2_CNN_BiLSTM":
-        return Model2CNNBiLSTM()
-    else:
-        raise ValueError(f"Unknown model name: {model_name}")
+
+import json, torch, torch.nn as nn, torch.nn.functional as F
+from pathlib import Path
+from transformers import AutoModel, AutoTokenizer
+from safetensors.torch import load_file as safe_load
+
+# -------- Base pooling --------
+class PoolingLayer(nn.Module):
+    def __init__(self, pooling="masked_mean"):
+        super().__init__()
+        assert pooling in ["cls","masked_mean","masked_max"]
+        self.pooling = pooling
+    def forward(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
+
+# -------- Model defs --------
+def _base_model(name):
+    return AutoModel.from_pretrained(name)
+
+class Model1_WCB(nn.Module):
+    def __init__(self, name, num_labels=2, dropout=0.3):
+        super().__init__()
+        self.bert = _base_model(name)
+        H = self.bert.config.hidden_size
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(H, num_labels)
+    def forward(self, ids, mask):
+        out = self.bert(input_ids=ids, attention_mask=mask)
+        cls = out.last_hidden_state[:,0,:]
+        return self.fc(self.dropout(cls))
+
+class Model2_WCB_BiLSTM(nn.Module):
+    def __init__(self, name, num_labels=2, hidden=128, dropout=0.3, pooling="masked_mean"):
+        super().__init__()
+        self.bert = _base_model(name)
+        H = self.bert.config.hidden_size
+        self.lstm = nn.LSTM(H, hidden, bidirectional=True, batch_first=True)
+        self.pool = PoolingLayer(pooling)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(hidden*2, num_labels)
+    def forward(self, ids, mask):
+        seq = self.bert(input_ids=ids, attention_mask=mask).last_hidden_state
+        x,_ = self.lstm(seq)
+        x = self.pool(x, mask)
+        return self.fc(self.dropout(x))
+
+class Model3_WCB_CNN_BiLSTM(nn.Module):
+    def __init__(self, name, num_labels=2, hidden=128, dropout=0.3, pooling="masked_mean"):
+        super().__init__()
+        self.bert = _base_model(name)
+        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.lstm = nn.LSTM(128, hidden, bidirectional=True, batch_first=True)
+        self.pool = PoolingLayer(pooling)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(hidden*2, num_labels)
+    def forward(self, ids, mask):
+        out = self.bert(input_ids=ids, attention_mask=mask).last_hidden_state
+        x = F.relu(self.c1(out.transpose(1,2)))
+        x = F.relu(self.c2(x)).transpose(1,2)
+        x,_ = self.lstm(x)
+        x = self.pool(x, mask)
+        return self.fc(self.dropout(x))
+
+class Model4_WCB_4Layer_BiLSTM(nn.Module):
+    def __init__(self, name, num_labels=2, hidden=128, dropout=0.3, pooling="masked_mean"):
+        super().__init__()
+        self.bert = _base_model(name)
+        H = self.bert.config.hidden_size
+        self.w = nn.Parameter(torch.ones(4))
+        self.lstm = nn.LSTM(H, hidden, bidirectional=True, batch_first=True)
+        self.pool = PoolingLayer(pooling)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(hidden*2, num_labels)
+    def _pool_layers(self, hs):
+        last4 = hs[-4:]
+        w = F.softmax(self.w, 0)
+        return sum(w[i]*last4[i] for i in range(4))
+    def forward(self, ids, mask):
+        out = self.bert(input_ids=ids, attention_mask=mask, output_hidden_states=True)
+        seq = self._pool_layers(out.hidden_states)
+        x,_ = self.lstm(seq)
+        x = self.pool(x, mask)
+        return self.fc(self.dropout(x))
+
+# -------- Factory & Loader --------
+def _build(arch, base_model, num_labels, pooling):
+    if arch == "WCB":
+        return Model1_WCB(base_model, num_labels)
+    if arch == "WCB_BiLSTM":
+        return Model2_WCB_BiLSTM(base_model, num_labels, pooling=pooling)
+    if arch == "WCB_CNN_BiLSTM":
+        return Model3_WCB_CNN_BiLSTM(base_model, num_labels, pooling=pooling)
+    if arch == "WCB_4Layer_BiLSTM":
+        return Model4_WCB_4Layer_BiLSTM(base_model, num_labels, pooling=pooling)
+    raise ValueError(f"Unknown architecture: {arch}")
+
+def load_model(model_dir: str):
+    """
+    โหลดโมเดลจากโฟลเดอร์โมเดล (ที่มี config.json + model.safetensors)
+    Return: tokenizer, model (eval mode), config(dict)
+    """
+    d = Path(model_dir)
+    cfg = json.loads((d/"config.json").read_text(encoding="utf-8"))
+    arch    = cfg.get("architecture","WCB")
+    base    = cfg.get("base_model","airesearch/wangchanberta-base-att-spm-uncased")
+    nlabel  = int(cfg.get("num_labels",2))
+    pooling = cfg.get("pooling_after_lstm","masked_mean")
+
+    model = _build(arch, base, nlabel, pooling)
+    sd = safe_load(str(d/"model.safetensors"))
+    model.load_state_dict(sd, strict=False)
+    model.eval()
+
+    tok = AutoTokenizer.from_pretrained(base, use_fast=True)
+    return tok, model, cfg