Update README.md

README.md

@@ -84,102 +84,51 @@ base_model:
 ---
 ## How to use the model
 ```
-import torch
-import torch.nn as nn
-import numpy as np
+import torch, json
 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)
+
+repo_or_dir = "LangQuant/LQ-Kbert-base" 
+texts = [
+    "비트코인 조정 후 반등, 투자심리 개선",
+    "환율 급등에 증시 변동성 확대",
+    "비트 그만 좀 내려라 진짜..",
+    "폭락ㅠㅠㅜㅠㅜ 다 팔아야할까요?"
+]
+
+
+tokenizer = AutoTokenizer.from_pretrained(repo_or_dir, local_files_only=True)
+model = AutoModel.from_pretrained(repo_or_dir, trust_remote_code=True, local_files_only=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device).eval()
+
+
+enc = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=200).to(device)
+with torch.inference_mode():
     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]
+IDX2SENTI = {0:"strong_pos",1:"weak_pos",2:"neutral",3:"weak_neg",4:"strong_neg"}
+IDX2ACT   = {0:"buy",1:"hold",2:"sell",3:"avoid",4:"info_only",5:"ask_info"}
+EMO_LIST  = ["greed","fear","confidence","doubt","anger","hope","sarcasm"]
+
 
-    emos = [EMO_LIST[i] for i,p in enumerate(emo_p) if p >= emo_threshold]
+for i, t in enumerate(texts):
+    senti = int(out["logits_senti"][i].argmax().item())
+    act   = int(out["logits_act"][i].argmax().item())
+    emo_p = torch.sigmoid(out["logits_emo"][i]).tolist()
+    reg   = torch.clamp(out["pred_reg"][i], 0, 1).tolist()
+    emos = [EMO_LIST[j] for j,p in enumerate(emo_p) if p >= 0.5]
 
-    return {
-        "text": text,
+    result = {
+        "text": t,
         "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)),
+        "pred_action_signal":      IDX2ACT[act],
+        "pred_emotions":           emos,
+        "pred_certainty":  float(reg[0]),
+        "pred_relevance":  float(reg[1]),
+        "pred_toxicity":   float(reg[2]),
     }
+    print(json.dumps(result, ensure_ascii=False))
 
-# ---- 메인 ----
-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}")
 ```
 ---