Create README.md

README.md

@@ -0,0 +1,69 @@
+---
+language:
+- ko
+base_model:
+- beomi/KcELECTRA-base
+pipeline_tag: text-classification
+tags:
+- emotion
+- sentiment
+---
+
+## Text의 감정을 분석하기 위한 모델입니다.
+
+### transformers로 모델을 받은 후 아래의 코드를 실행, sample_text에 문장을 넣으면 분석해줍니다.
+
+```
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch.nn.functional as F
+import os
+
+# 모델이랑 토크나이저 불러오기.
+MODEL_DIR = os.path.join(OUTPUT_DIR, "best_model")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
+model.eval().to("cuda" if torch.cuda.is_available() else "cpu")
+
+# 감정 카테고리 - 맨 위에 저장해놓은 그거
+CATEGORIES = ["happy", "embarrass", "anger", "unrest", "damaged", "sadness"]
+
+# 추론 함수. 여기에 txt를 넣으면 결과가 나옴.
+def predict_emotion(texts):
+    if isinstance(texts, str):
+        texts = [texts]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    inputs = tokenizer(
+        texts,
+        padding=True,
+        truncation=True,
+        max_length=64,
+        return_tensors="pt"
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = F.softmax(outputs.logits, dim=-1)
+        preds = probs.argmax(dim=-1)
+
+    results = []
+    for i, text in enumerate(texts):
+        results.append({
+            "text": text,
+            "pred_label": CATEGORIES[preds[i].item()],
+            "probabilities": {CATEGORIES[j]: round(probs[i][j].item(), 4) for j in range(len(CATEGORIES))}
+        })
+    return results
+
+# 넣을 문장 예시. 나중에는 사용자의 다이어리 내용이 들어갈 것.
+sample_texts = [
+    "나 오늘 기뻐서 울었어"
+]
+
+for r in predict_emotion(sample_texts):
+    print(f"문장: {r['text']}")
+    print(f"예측 감정: {r['pred_label']}")
+    print("확률 분포:", r["probabilities"])
+    print("-" * 60)
+```