koelectra-emotion-7-emotion-base-onnx

koelectra-emotion-7-emotion-base (ONNX)

한국어 문장의 감정을 분류하는 KoELECTRA-base-v3 파인튜닝 모델의 ONNX 변환본입니다. 문장 1개를 입력하면 7개 클래스 — 감정 6종(기쁨, 분노, 상처, 불안, 당황, 슬픔) + 중립(감정 없음) — 의 확률을 출력합니다.

PyTorch·transformers·CUDA가 전혀 필요 없습니다. 런타임 의존성은 onnxruntime + tokenizers + numpy 3개뿐이며, CPU에서 문장 1개당 수십 ms로 동작해 GPU 없는 배포 환경에 적합합니다.

• Base model: monologg/koelectra-base-v3-discriminator
• 원본(PyTorch) 모델: MelissaJ/koelectra-emotion-7-emotion-base
• 입력: 한국어 단문(발화 1개), max_length 128 토큰
• 변환: opset 17, fp32, 동적 배치/시퀀스 축 (torch.onnx.export)
• 변환 검증: 원본 torch 로짓과 1e-3 이내 일치 (parity 테스트 통과)

파일 구성

설치

pip install onnxruntime tokenizers numpy huggingface_hub

사용법

import json
from pathlib import Path

import numpy as np
import onnxruntime as ort
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer

model_dir = Path(snapshot_download("MelissaJ/koelectra-emotion-7-emotion-base-onnx"))

tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
tokenizer.enable_truncation(128)
session = ort.InferenceSession(
    str(model_dir / "model.onnx"), providers=["CPUExecutionProvider"]
)
config = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
id2label = {int(i): label for i, label in config["id2label"].items()}


def classify_emotion(text: str) -> dict[str, float]:
    enc = tokenizer.encode(text)
    feeds = {
        "input_ids": np.array([enc.ids], dtype=np.int64),
        "attention_mask": np.array([enc.attention_mask], dtype=np.int64),
        "token_type_ids": np.array([enc.type_ids], dtype=np.int64),
    }
    logits = session.run(None, feeds)[0][0]
    exp = np.exp(logits - logits.max())
    probs = exp / exp.sum()
    return dict(
        sorted(
            ((id2label[i], float(p)) for i, p in enumerate(probs)),
            key=lambda x: -x[1],
        )
    )


print(classify_emotion("오늘 드디어 합격 소식을 들었어!"))
# {'기쁨': 0.998, ...}

출력 예시

"오늘 드디어 합격 소식을 들었어!"      → 기쁨 0.998
"밤마다 미래가 걱정돼서 잠이 안 와."   → 불안 0.489 (2위 상처 0.300)
"오늘 회의는 3시에 시작합니다."        → 중립 1.000
"점심 뭐 먹을까?"                      → 중립 1.000

GPU 가속 (선택)

CPU만으로 충분히 빠르지만, 필요하면 onnxruntime-gpu 설치 후 providers만 바꾸면 됩니다:

session = ort.InferenceSession(
    str(model_dir / "model.onnx"),
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

학습 데이터

총 train 169,430 / val 20,935 문장. 클래스 불균형(기쁨이 적음)은 weighted cross-entropy로 보정했습니다.

대화 스타일 증강 (v2, 2026-08): train/val 각 50%에 호명("○○님!", "○○아,"), 구어 감탄("아이고, 진짜!", "어휴," 등), 말줄임("...")을 라벨 불변으로 전 클래스에 균일 주입했습니다. v1은 중립 소스가 전부 뉴스·챗봇체라 "호명·구어 감탄 = 중립"이라는 도메인 지름길을 학습했고, 캐릭터 대화체 감정 발화("멜리사님! 정말 화가 난다!")가 중립 확률 ~1.0으로 포화되는 문제가 실측됐습니다. 균일 주입으로 이 지름길을 차단했습니다.

학습 설정

• 8 epochs, batch 32, lr 3e-5, max_length 128, fp16, seed 42
• warmup 10%, weight decay 0.01, f1_macro 기준 best model 선택 (early stopping patience 2)

성능 (val 20,935문장, 스타일 증강 적용 분포)

ONNX 변환은 무손실에 가깝습니다(로짓 차이 1e-3 이내, v2 실측 최대 4e-5). 아래는 원본 모델 기준 성능이며, v2의 val은 증강 분포라 v1 수치(acc 0.6790 / f1_macro 0.6744)와 직접 비교할 수 없습니다.

호명·감탄 회귀 프로브(v1에서 전부 중립 ~1.0 포화)가 v2에서 올바른 감정으로 복귀: "멜리사님! 정말 화가 난다!" → 분노 0.79, "...아이고, 진짜! ○○님는 나한테 맨날 일만 시키지!" → 분노 0.98. 호명 섞인 중립("○○님, 오늘 회의는 3시에 시작합니다")은 중립 유지.

한계

• 상처/슬픔/불안 경계가 흐립니다. 감성대화말뭉치 특성상 라벨러 간에도 갈리는 부류라 f1이 0.5 안팎입니다. 확신 임계값(최상위 확률 0.5 미만이면 판단 보류) 병행을 권장합니다.
• 반말 구어 평서문 중립("부산까지 세 시간 걸려")은 감정(주로 기쁨)으로 오분류될 수 있습니다. 감정 학습 데이터가 대부분 반말 구어체 발화라서 생긴 편향입니다. 뉴스형 문장, 일상 질문, 경어체/문어체 평서문 중립은 잘 잡힙니다.
• 단문(발화 1개) 기준으로 학습되어, 긴 문단보다 문장 단위 입력이 학습 분포와 맞습니다.