noridorimari/onnx_emotion_classifier

문장의 감정을 분석하기 위한 모델입니다.

KcELECTRA 모델을 기반으로 약 22만개의 감정 문장을 학습하였습니다.

기존 FP16 모델을 ONNX FP16으로 변환하여 CPU 환경에서의 추론 속도를 약 2배 단축하였습니다.

감정은 총 6개의 카테고리로 도출되며 기쁨, 당황, 분노, 불안, 상처, 슬픔 입니다.

# onnx로 바꾸면서 softmax가 풀렸으므로 다시 numpy를 사용해 만들어줌
def softmax(x):
    x = np.array(x)
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=-1, keepdims=True)

def predict(text: str):
    text = clean(text)

    inputs = tokenizer(text, return_tensors="np")

    ort_inputs = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"]
    }

    logits = session.run(["logits"], ort_inputs)[0]  # (1, 6)
    probs = softmax(logits)[0]  # shape: (6,)

    result = {id2label[i]: float(probs[i]) for i in range(6)}
    return result

# pip install optimum[onnxruntime]

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

model_id = "noridorimari/onnx_emotion_classifier"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForSequenceClassification.from_pretrained(model_id)

inputs = tokenizer("오늘 기분이 별로 좋지 않아.", return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits)

@misc{lee2021kcelectra,
  author = {Junbum Lee},
  title = {KcELECTRA: Korean comments ELECTRA},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/Beomi/KcELECTRA}}
}