File size: 1,501 Bytes
136c55f
e850536
 
 
 
 
d2df06f
 
 
 
e850536
136c55f
d2df06f
 
 
 
 
e850536
d2df06f
 
 
 
e850536
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
# emotion_engine.py (์ˆ˜์ • ํ›„ ์ตœ์ข… ๋ฒ„์ „)

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import os

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import os

def load_emotion_classifier():
    # --- ์ด ๋ถ€๋ถ„์„ ์ˆ˜์ •ํ•ฉ๋‹ˆ๋‹ค ---
    # ๋กœ์ปฌ ๊ฒฝ๋กœ ๋Œ€์‹ , Hugging Face Hub์˜ ๋ชจ๋ธ ID๋ฅผ ์‚ฌ์šฉํ•ฉ๋‹ˆ๋‹ค.
    MODEL_PATH = "koons/korean-emotion-classifier-final" # "์‚ฌ์šฉ์ž์ด๋ฆ„/๋ชจ๋ธ์ด๋ฆ„" ํ˜•์‹

    print(f"Hugging Face Hub ๋ชจ๋ธ '{MODEL_PATH}'์—์„œ ๋ชจ๋ธ์„ ๋ถˆ๋Ÿฌ์˜ต๋‹ˆ๋‹ค...")

    try:
        # local_files_only ์˜ต์…˜์„ ์ œ๊ฑฐํ•˜์—ฌ ์˜จ๋ผ์ธ์—์„œ ๋‹ค์šด๋กœ๋“œํ•˜๋„๋ก ํ•ฉ๋‹ˆ๋‹ค.
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
        print("โœ… Hugging Face Hub ๋ชจ๋ธ ๋กœ๋”ฉ ์„ฑ๊ณต!")

    except Exception as e:
        print(f"โŒ ๋ชจ๋ธ ๋กœ๋”ฉ ์ค‘ ์˜ค๋ฅ˜: {e}")
        return None
    
    device = 0 if torch.cuda.is_available() else -1
    emotion_classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
    
    return emotion_classifier

# predict_emotion ํ•จ์ˆ˜๋Š” ๊ทธ๋Œ€๋กœ ๋‘ก๋‹ˆ๋‹ค.
def predict_emotion(classifier, text):
    if not text or not text.strip(): return "๋‚ด์šฉ ์—†์Œ"
    if classifier is None: return "์˜ค๋ฅ˜: ๊ฐ์ • ๋ถ„์„ ์—”์ง„ ์ค€๋น„ ์•ˆ๋จ."
    result = classifier(text)
    return result[0]['label']