Upload inference/pipeline.py with huggingface_hub

inference/pipeline.py

@@ -0,0 +1,265 @@
+# inference/pipeline.py
+import os
+import json
+import torch
+import yaml
+from typing import List, Dict, Any, Optional
+
+from tokenizer.thai_tokenizer import ThaiTokenizer
+from model.encoder import ThaiTransformerEncoder, ModelConfig
+from model.heads.ner_head import NERHead
+from model.heads.sentiment_head import SentimentHead
+from model.heads.qa_head import QAHead
+
+
+# label maps สำหรับ decode output กลับเป็น string
+NER_ID2LABEL = {
+    0: "O",
+    1: "B-PERSON", 2: "I-PERSON",
+    3: "B-ORGANIZATION",  4: "I-ORGANIZATION",
+    5: "B-LOCATION",  6: "I-LOCATION",
+}
+SENTIMENT_ID2LABEL = {0: "negative", 1: "neutral", 2: "positive"}
+
+
+class ThaiNLPModel(torch.nn.Module):
+    """รวม encoder + 3 heads เป็น module เดียว สำหรับ load/save"""
+    def __init__(self, config: ModelConfig, num_ner_labels: int = 7):
+        super().__init__()
+        self.encoder        = ThaiTransformerEncoder(config)
+        self.ner_head       = NERHead(config.d_model, num_ner_labels)
+        self.sentiment_head = SentimentHead(config.d_model, num_classes=3)
+        self.qa_head        = QAHead(config.d_model)
+
+
+class ThaiNLPPipeline:
+    """
+    High-level inference class
+    โหลด model ครั้งเดียวแล้วเรียก predict() ได้เรื่อยๆ
+    """
+
+    def __init__(self, model_dir: str, device: str = "auto", checkpoint_name: str = "checkpoint_best"):
+        # ── Device ───────────────────────────────────────────────────────
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+
+        # ── Load config ───────────────────────────────────────────────────
+        config_path = os.path.join(model_dir, "config.yaml")
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"ไม่พบ config.yaml ใน {model_dir}")
+        with open(config_path) as f:
+            raw_config = yaml.safe_load(f)
+        model_cfg = ModelConfig(**raw_config["model"])
+
+        # ── Load tokenizer ────────────────────────────────────────────────
+        self.tokenizer = ThaiTokenizer.from_pretrained(
+            os.path.join(model_dir, "tokenizer")
+        )
+
+        # ── Load model ────────────────────────────────────────────────────
+        self.model = ThaiNLPModel(model_cfg)
+        ckpt_path  = os.path.join(model_dir, checkpoint_name, "checkpoint.pt")
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(f"ไม่พบ checkpoint: {ckpt_path}")
+
+        ckpt = torch.load(ckpt_path, map_location=self.device)
+        self.model.load_state_dict(ckpt["model_state"])
+        self.model.to(self.device)
+        self.model.eval()
+
+        step = ckpt.get("global_step", "unknown")
+        metric = ckpt.get("best_metric", "unknown")
+        print(f"pipeline ready on {self.device} (loaded {checkpoint_name} from step {step} with best_val_loss={metric})")
+
+    # ─────────────────────────────────────────────────────────────────────
+    # predict — entry point หลัก
+    # ─────────────────────────────────────────────────────────────────────
+
+    def predict(
+        self,
+        text:     str,
+        tasks:    List[str],
+        question: Optional[str] = None,   # ต้องการสำหรับ QA
+        context:  Optional[str] = None,   # ต้องการสำหรับ QA
+    ) -> Dict[str, Any]:
+        """
+        Parameters
+        ----------
+        text     : input text สำหรับ NER และ Sentiment
+        tasks    : list ของ task ที่ต้องการ ["ner", "sentiment", "qa"]
+        question : question string (เฉพาะ QA)
+        context  : context string (เฉพาะ QA)
+
+        Returns
+        -------
+        dict ที่มี key ตาม tasks ที่ขอ
+        """
+        results = {}
+
+        with torch.no_grad():
+            # ── NER ──────────────────────────────────────────────────────
+            if "ner" in tasks:
+                results["ner"] = self._predict_ner(text)
+
+            # ── Sentiment ────────────────────────────────────────────────
+            if "sentiment" in tasks:
+                results["sentiment"] = self._predict_sentiment(text)
+
+            # ── QA ───────────────────────────────────────────────────────
+            if "qa" in tasks:
+                if question is None or context is None:
+                    results["qa"] = {"error": "QA ต้องการ question และ context"}
+                else:
+                    results["qa"] = self._predict_qa(question, context)
+
+        return results
+
+    # ─────────────────────────────────────────────────────────────────────
+    # Task-specific predict methods
+    # ─────────────────────────────────────────────────────────────────────
+
+    def _encode(self, input_ids, attention_mask):
+        """Shared encoder forward"""
+        ids  = torch.tensor([input_ids],      dtype=torch.long).to(self.device)
+        mask = torch.tensor([attention_mask], dtype=torch.long).to(self.device)
+        hidden, _ = self.model.encoder(ids, mask)
+        return hidden, mask
+
+    def _predict_ner(self, text: str) -> List[Dict[str, str]]:
+        """
+        คืน list ของ {"token": str, "label": str}
+        กรอง [CLS], [SEP], padding ออก และ merge subwords กลับเป็นคำ
+        """
+        encoded    = self.tokenizer.batch_encode(
+            [text], max_length=512, padding=False, return_tensors=False
+        )
+        input_ids  = encoded["input_ids"][0]
+        attn_mask  = encoded["attention_mask"][0]
+
+        hidden, _  = self._encode(input_ids, attn_mask)
+        logits     = self.model.ner_head(hidden)          # (1, T, num_labels)
+        pred_ids   = logits[0].argmax(dim=-1).tolist()    # (T,)
+
+        # decode tokens กลับเป็น string แล้ว zip กับ label
+        pieces = self.tokenizer.sp.id_to_piece(input_ids)
+        special = {
+            self.tokenizer.cls_id,
+            self.tokenizer.sep_id,
+            self.tokenizer.pad_id,
+        }
+
+        entities = []
+        current_word  = ""
+        current_label = "O"
+
+        for token_id, label_id in zip(input_ids, pred_ids):
+            if token_id in special:
+                continue
+
+            piece = self.tokenizer.sp.id_to_piece([token_id])[0]
+            label = NER_ID2LABEL.get(label_id, "O")
+
+            # SentencePiece ใช้ "▁" นำหน้า subword แรกของคำ
+            if piece.startswith("▁") or not current_word:
+                # บันทึกคำก่อนหน้า (ถ้ามี)
+                if current_word:
+                    entities.append({
+                        "token": current_word,
+                        "label": current_label,
+                    })
+                current_word  = piece.lstrip("▁")
+                current_label = label
+            else:
+                # subword ต่อเนื่อง — merge เข้ากับคำปัจจุบัน
+                current_word += piece
+
+        # บันทึกคำสุดท้าย
+        if current_word:
+            entities.append({"token": current_word, "label": current_label})
+
+        return entities
+
+    def _predict_sentiment(self, text: str) -> Dict[str, Any]:
+        """
+        คืน {"label": str, "confidence": float, "scores": dict}
+        """
+        encoded   = self.tokenizer.batch_encode(
+            [text], max_length=512, padding=False, return_tensors=False
+        )
+        hidden, mask = self._encode(
+            encoded["input_ids"][0],
+            encoded["attention_mask"][0],
+        )
+        mask_tensor = torch.tensor(
+            [encoded["attention_mask"][0]], dtype=torch.long
+        ).to(self.device)
+
+        logits = self.model.sentiment_head(hidden, mask_tensor)  # (1, 3)
+        probs  = logits.softmax(dim=-1)[0].tolist()
+
+        pred_id     = int(logits.argmax(dim=-1).item())
+        pred_label  = SENTIMENT_ID2LABEL[pred_id]
+        confidence  = round(probs[pred_id], 4)
+
+        return {
+            "label":      pred_label,
+            "confidence": confidence,
+            "scores": {
+                SENTIMENT_ID2LABEL[i]: round(p, 4)
+                for i, p in enumerate(probs)
+            },
+        }
+
+    def _predict_qa(self, question: str, context: str) -> Dict[str, Any]:
+        """
+        คืน {"answer": str, "start": int, "end": int, "confidence": float}
+        """
+        encoded       = self.tokenizer.encode_qa(question, context, max_length=512)
+        input_ids     = encoded["input_ids"]
+        attn_mask     = encoded["attention_mask"]
+        context_start = encoded["context_start"]
+
+        hidden, _ = self._encode(input_ids, attn_mask)
+
+        start_logits, end_logits = self.model.qa_head(
+            hidden, context_start=context_start
+        )   # (1, T) each
+
+        # หา (start, end) ที่ให้ score สูงสุดโดย start ≤ end
+        start_logits = start_logits[0]   # (T,)
+        end_logits   = end_logits[0]     # (T,)
+        seq_len      = len(input_ids)
+
+        best_score = float("-inf")
+        best_start = context_start
+        best_end   = context_start
+
+        # จำกัด span ไม่เกิน 30 tokens (คำตอบยาวเกินนี้ไม่สมเหตุสมผล)
+        MAX_ANSWER_LEN = 30
+
+        for s in range(context_start, seq_len):
+            for e in range(s, min(s + MAX_ANSWER_LEN, seq_len)):
+                score = start_logits[s].item() + end_logits[e].item()
+                if score > best_score:
+                    best_score = score
+                    best_start = s
+                    best_end   = e
+
+        # decode answer กลับเป็น string
+        answer_ids = input_ids[best_start:best_end + 1]
+        answer     = self.tokenizer.decode(answer_ids, skip_special_tokens=True)
+
+        # confidence = softmax score ของ best span (normalize คร่าวๆ)
+        start_probs = start_logits.softmax(dim=-1)
+        end_probs   = end_logits.softmax(dim=-1)
+        confidence  = round(
+            (start_probs[best_start] * end_probs[best_end]).item(), 4
+        )
+
+        return {
+            "answer":     answer,
+            "start":      best_start,
+            "end":        best_end,
+            "confidence": confidence,
+        }