VERITAS 부착: 답+마지막층 은닉 동시 반환 핸들러

handler.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+"""VERITAS 부착 핸들러 — JGOS-398B-fp8 (9-팀장 2026-09-17)
+
+무엇을 하나
+  한 번의 호출로 **답과 속 신호를 같이** 돌려준다.
+  속 신호 = 마지막 층·마지막 토큰의 은닉벡터(가공 안 된 원본값).
+  이것이 VERITAS 가 읽는 바로 그 값이며, **토큰을 한 개도 더 쓰지 않는다**
+  (이미 계산된 것을 꺼내 올 뿐이다).
+
+두 가지 모드
+  mode="generate" : 답 생성 + 그 답까지 포함한 은닉 (문항 풀이용)
+  mode="embed"    : 주어진 글자열의 은닉만 (이미 뽑아 둔 답을 채점할 때)
+
+🔴 규율
+  - 잘림(length)을 결과에 반드시 표시한다. 자원 소진을 오답으로 세지 않기 위해서다.
+  - 은닉은 정규화하지 않는다. 정규화하면 VERITAS 가 쓰는 크기 정보가 날아간다.
+  - 생성분이 reasoning 쪽으로 빠지는 템플릿이면 그것까지 합쳐 돌려준다(버리지 않는다).
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        self.tok = AutoTokenizer.from_pretrained(path)
+        if self.tok.pad_token is None:
+            self.tok.pad_token = self.tok.eos_token
+        self.tok.padding_side = "left"          # 🔴 디코더 모델 배치의 필수 조건
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path, dtype="auto", device_map="auto"
+        ).eval()
+        cfg = self.model.config
+        cfg = cfg.get_text_config() if hasattr(cfg, "get_text_config") else cfg
+        self.n_layers = cfg.num_hidden_layers
+
+    # ---- 내부: 마지막 층·마지막 토큰 은닉 ------------------------------
+    @torch.no_grad()
+    def _hidden(self, ids) -> List[List[float]]:
+        out = self.model(**ids, output_hidden_states=True, use_cache=False)
+        h = out.hidden_states[-1]                       # 마지막 층
+        mask = ids["attention_mask"]
+        last = mask.sum(dim=1) - 1                      # 좌패딩이라 실제 마지막 위치
+        picked = h[torch.arange(h.size(0), device=h.device), last, :]
+        return picked.float().cpu().tolist()
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        inputs = data.get("inputs")
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        p = data.get("parameters") or {}
+        mode = p.get("mode", "generate")
+
+        # ── 은닉만 ──────────────────────────────────────────────────
+        if mode == "embed":
+            ids = self.tok(inputs, return_tensors="pt", padding=True,
+                           truncation=True, max_length=p.get("max_length", 32768))
+            ids = {k: v.to(self.model.device) for k, v in ids.items()}
+            return {"hidden": self._hidden(ids), "layer": self.n_layers,
+                    "normalized": False}
+
+        # ── 생성 + 은닉 ─────────────────────────────────────────────
+        max_new = int(p.get("max_new_tokens", 16384))
+        n = int(p.get("n", 1))
+        temp = float(p.get("temperature", 0.0))
+        texts = []
+        for s in inputs:
+            if p.get("chat", True):
+                msgs = [{"role": "user", "content": s}]
+                kw = {}
+                if p.get("enable_thinking") is not None:
+                    kw["enable_thinking"] = p["enable_thinking"]
+                try:
+                    texts.append(self.tok.apply_chat_template(
+                        msgs, tokenize=False, add_generation_prompt=True, **kw))
+                except TypeError:
+                    texts.append(self.tok.apply_chat_template(
+                        msgs, tokenize=False, add_generation_prompt=True))
+            else:
+                texts.append(s)
+        enc = self.tok(texts, return_tensors="pt", padding=True)
+        enc = {k: v.to(self.model.device) for k, v in enc.items()}
+        gen_kw = dict(max_new_tokens=max_new, pad_token_id=self.tok.pad_token_id)
+        if temp > 0:
+            gen_kw.update(do_sample=True, temperature=temp, num_return_sequences=n)
+        else:
+            gen_kw.update(do_sample=False)
+        with torch.no_grad():
+            g = self.model.generate(**enc, **gen_kw)
+        plen = enc["input_ids"].shape[1]
+        new = g[:, plen:]
+        outs = self.tok.batch_decode(new, skip_special_tokens=True)
+        ntok = (new != self.tok.pad_token_id).sum(dim=1).tolist()
+        # 🔴 잘림 표시 — 자원 소진을 오답으로 세지 않기 위해 결과에 남긴다
+        hit_cap = [int(t >= max_new - 2) for t in ntok]
+
+        att = torch.ones_like(g)
+        att[g == self.tok.pad_token_id] = 0
+        hid = self._hidden({"input_ids": g, "attention_mask": att})
+        return {"generated": outs, "hidden": hid, "n_tokens": ntok,
+                "hit_cap": hit_cap, "layer": self.n_layers, "normalized": False}