Update README.md

README.md

@@ -52,23 +52,68 @@ KoELECTRA 기반의 한국어(특히 일기/심리 기록) 감정 분류 모델
 ## How to Get Started with the Model
 ```python
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch.nn.functional as F
 import torch
+import torch.nn.functional as F
+
+# 1) Load Model & Tokenizer
+MODEL_NAME = "LimYeri/HowRU-KoELECTRA-Emotion-Classifier"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+
+# GPU 사용 가능 시 자동 전환
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
+
+# 감정 라벨 매핑 (id2label)
+id2label = model.config.id2label
 
-model_name = "LimYeri/HowRU-KoELECTRA-Emotion-Classifier"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# 2) Inference Function
+def predict_emotion(text: str):
+    """
+    Returns:
+        - top1_pred: 예측된 감정 라벨
+        - probs_sorted: 감정별 확률(내림차순)
+        - top2_pred: 상위 두 개의 감정
+    """
 
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True)
+    # 토크나이징
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    ).to(device)
+
+    # 추론
     with torch.no_grad():
         logits = model(**inputs).logits
         probs = F.softmax(logits, dim=-1)[0]
-    label_id = probs.argmax().item()
-    return label_id, probs.tolist()
 
-print(predict("오늘 정말 기분이 좋고 행복한 하루였어!"))
+    # 정렬된 확률
+    probs_sorted = sorted(
+        [(id2label[i], float(probs[i])) for i in range(len(probs))],
+        key=lambda x: x[1],
+        reverse=True
+    )
+
+    top1_pred = probs_sorted[0]
+    top2_pred = probs_sorted[:2]
+
+    return {
+        "text": text,
+        "top1_emotion": top1_pred,
+        "top2_emotions": top2_pred,
+        "all_probabilities": probs_sorted,
+    }
+
+
+# 3) Example
+result = predict_emotion("오늘 정말 기분이 좋고 행복한 하루였어!")
+print(result)
 ```
 
 ---