Update README.md

README.md

@@ -81,6 +81,106 @@ base_model:
   - 매수/매도/보유/회피/질문/정보.  
 
 
+---
+## How to use the model
+```
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+from huggingface_hub import hf_hub_download
+
+# ---- 상수 정의 ----
+REPO_ID = "langquant/LQ-Kbert-base"
+CKPT_RELPATH = "model/lq-kbert-base.pt"
+
+SENTI_MAP = {'strong_pos':0,'weak_pos':1,'neutral':2,'weak_neg':3,'strong_neg':4}
+ACT_MAP   = {'buy':0,'hold':1,'sell':2,'avoid':3,'info_only':4,'ask_info':5}
+EMO_LIST  = ['greed','fear','confidence','doubt','anger','hope','sarcasm']
+IDX2SENTI = {v:k for k,v in SENTI_MAP.items()}
+IDX2ACT   = {v:k for k,v in ACT_MAP.items()}
+
+def sigmoid(x): return 1/(1+np.exp(-x))
+
+# ---- 모델 정의 ----
+class KbertMTL(nn.Module):
+    def __init__(self, base_model, hidden=768):
+        super().__init__()
+        self.bert = base_model
+        self.head_senti = nn.Linear(hidden, 5)
+        self.head_act   = nn.Linear(hidden, 6)
+        self.head_emo   = nn.Linear(hidden, 7)
+        self.head_reg   = nn.Linear(hidden, 3)
+        self.has_token_type = getattr(self.bert.embeddings, "token_type_embeddings", None) is not None
+
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
+        kwargs = dict(input_ids=input_ids, attention_mask=attention_mask)
+        if self.has_token_type and token_type_ids is not None:
+            kwargs["token_type_ids"] = token_type_ids
+        out = self.bert(**kwargs)
+        h = out.last_hidden_state[:, 0]  # [CLS]
+        return {
+            "logits_senti": self.head_senti(h),
+            "logits_act":   self.head_act(h),
+            "logits_emo":   self.head_emo(h),
+            "pred_reg":     self.head_reg(h)
+        }
+
+# ---- 체크포인트 로드 ----
+def load_ckpt_from_hub():
+    path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_RELPATH)
+    obj = torch.load(path, map_location="cpu")
+    return obj
+
+# ---- 모델 및 토크나이저 구성 ----
+def build_model_and_tokenizer(ckpt_obj, hidden=768):
+    model_name = ckpt_obj.get("model_name", "klue/bert-base")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    base = AutoModel.from_pretrained(model_name)
+    model = KbertMTL(base_model=base, hidden=hidden)
+    state_dict = ckpt_obj["state_dict"] if "state_dict" in ckpt_obj else ckpt_obj
+    model.load_state_dict(state_dict, strict=False)
+    emo_thr = float(ckpt_obj.get("emo_threshold", 0.5))
+    return model, tokenizer, emo_thr
+
+# ---- 추론 ----
+@torch.no_grad()
+def predict(text, model, tokenizer, device="cpu", max_len=200, emo_threshold=0.5):
+    model.to(device).eval()
+    enc = tokenizer([text], padding=True, truncation=True, max_length=max_len, return_tensors="pt").to(device)
+    out = model(**enc)
+
+    senti = out["logits_senti"].argmax(-1).item()
+    act   = out["logits_act"].argmax(-1).item()
+    emo_p = sigmoid(out["logits_emo"].cpu().numpy())[0]
+    reg   = out["pred_reg"].cpu().numpy()[0]
+
+    emos = [EMO_LIST[i] for i,p in enumerate(emo_p) if p >= emo_threshold]
+
+    return {
+        "text": text,
+        "pred_sentiment_strength": IDX2SENTI[senti],
+        "pred_action_signal": IDX2ACT[act],
+        "pred_emotions": emos,
+        "pred_certainty":  float(np.clip(reg[0], 0, 1)),
+        "pred_relevance":  float(np.clip(reg[1], 0, 1)),
+        "pred_toxicity":   float(np.clip(reg[2], 0, 1)),
+    }
+
+# ---- 메인 ----
+if __name__ == "__main__":
+    text = input("분석할 문장을 입력하세요: ").strip()
+    print("[모델 로드 중...]")
+    ckpt = load_ckpt_from_hub()
+    model, tokenizer, emo_thr = build_model_and_tokenizer(ckpt)
+
+    print("[추론 중...]")
+    result = predict(text, model, tokenizer, device="cuda" if torch.cuda.is_available() else "cpu", emo_threshold=emo_thr)
+
+    print("\n=== 결과 ===")
+    for k,v in result.items():
+        print(f"{k}: {v}")
+```
 ---
 
 ### 예시