koelectra-search-7-base-onnx

한국어 검색 라우터 (KoELECTRA Intent Router, ONNX)

한국어 사용자 발화를 분석해 7가지 기능 중 하나로 라우팅하는 의도 분류기입니다. LLM 챗봇 앞단에서 "이 발화를 이미지 검색으로 보낼지, 웹 검색으로 보낼지, 그냥 LLM이 답하게 둘지"를 매 턴 결정하는 용도로 설계했습니다.

monologg/koelectra-base-v3-discriminator를 파인튜닝한 뒤 ONNX로 변환했습니다. torch·transformers 불필요 — 런타임 의존성은 onnxruntime + tokenizers + numpy 3개뿐이며, CPU에서 단건 추론 평균 9.3ms (p99 10.9ms)로 동작합니다.

라벨 (7-class, 단일 라벨)

분류 기준

주제가 아니라 행위로 가릅니다. 같은 주제라도 요구하는 행위에 따라 라벨이 달라집니다.

블랙홀이 어떻게 생기는지 설명해줘          → chat    (LLM 지식으로 충분)
블랙홀 최신 관측 결과 인터넷에서 검색해줘   → text    (실시간 정보)
블랙홀 사진 보여줘                        → images
블랙홀 다큐 틀어줘                        → videos

URL이 있다고 무조건 extract가 아닙니다. 본문 추출을 요구하는지가 기준입니다.

https://example.com/gallery 이거랑 비슷한 사진 찾아줘  → images
이 링크 영상인데 재생해줘                             → videos
https랑 http 차이가 뭐야?                             → chat
이 링크 본문 긁어와                                   → extract

부정 표현은 마지막 지시를 따릅니다.

검색하지 말고 네가 아는대로 설명해줘     → chat
네 지식 말고 인터넷에서 검색해서 알려줘  → text
영상... 아니 사진으로 보여줘             → images

한국어 다의어 케이스를 학습 데이터에 의도적으로 포함해 키워드 지름길을 차단했습니다.

보도블록 사진 보여줘      → images  (보도 ≠ 언론보도)
커피 추출 원리 설명해줘   → chat    (추출 ≠ 본문추출)
에어컨 좀 틀어줘          → chat    (틀다 ≠ 영상재생)
회사에서 짤렸어           → chat    (짤 ≠ 짤방)

설치

pip install onnxruntime tokenizers numpy huggingface_hub

torch, transformers, CUDA는 필요 없습니다.

사용법

최소 예제

import json

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

repo = "MelissaJ/koelectra-search-7-base-onnx"
model_path = hf_hub_download(repo, "model.onnx")
tokenizer_path = hf_hub_download(repo, "tokenizer.json")
config_path = hf_hub_download(repo, "config.json")

tokenizer = Tokenizer.from_file(tokenizer_path)
tokenizer.enable_truncation(64)  # 학습 max_length와 동일 — 바꾸지 말 것
session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
with open(config_path, encoding="utf-8") as f:
    id2label = {int(k): v for k, v in json.load(f)["id2label"].items()}


def route(text: str) -> tuple[str, float]:
    enc = tokenizer.encode(text)
    logits = session.run(None, {
        "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),
    })[0][0]
    probs = np.exp(logits - logits.max())
    probs /= probs.sum()
    idx = int(probs.argmax())
    return id2label[idx], float(probs[idx])


print(route("지금 원달러 환율 얼마야"))        # ('text', 0.99...)
print(route("무라카미 하루키 소설 추천해줘"))   # ('books', 0.99...)
print(route("에어컨 좀 틀어줘"))               # ('chat', 0.99...)

노이즈 가드 (권장)

ㅋㅋ, ?, ..., ㅁㄴㅇㄹ 같은 무의미 입력은 모델에 태우지 말고 규칙으로 chat 폴백하는 것을 권장합니다. 학습 데이터에 이런 노이즈가 없기 때문입니다.

import re

_CORE = re.compile(r"[가-힣a-zA-Z0-9]")

def is_noise(text: str) -> bool:
    """완성형 한글/영숫자가 하나도 없으면 무의미 입력으로 간주."""
    return not _CORE.search(text.strip())

def route_safe(text: str) -> tuple[str, float]:
    if is_noise(text):
        return "chat", 1.0
    return route(text)

주의: 길이 기준(len < 3 등)으로 거르면 안 됩니다. 환율, 날씨, 주가 같은 정상적인 2글자 질의까지 차단됩니다.

완성형 라우터 클래스

배치 처리·top-k 확률·폴백 옵션이 포함된 OnnxIntentRouter 클래스와 CLI(--interactive, --sanity, --benchmark 등)는 저장소의 predict_onnx.py를 참고하세요.

from predict_onnx import OnnxIntentRouter

router = OnnxIntentRouter(model_dir="./koelectra-router-onnx")
result = router.predict("손흥민 어제 경기 하이라이트 틀어줘")[0]
print(result["label"], result["confidence"])   # videos 0.99...
print(result["probs"])                          # 전체 라벨 확률 (내림차순)

성능

주의: 위 test 점수를 실사용 성능으로 해석하면 안 됩니다. 학습·평가 데이터가 모두 동일 분포의 합성 데이터이므로, 실제 사용자 발화에서는 이보다 낮은 성능이 나올 수 있습니다.

학습 정보

• 베이스: monologg/koelectra-base-v3-discriminator (hidden 768, 12 layers)
• 데이터: 합성 한국어 발화 4,424건 (train 3,537 / val 444 / test 443, 층화 분할)
• 하이퍼파라미터: lr 3e-5, batch 32, 10 epochs, warmup 10%, weight decay 0.01, max_length 64, class weights로 불균형 보정 (최대 1.74:1)
• 검증: 시드 3개(42/43/44) 평균 val f1_macro 0.991 ± 0.002
• ONNX: opset 기본값, CPUExecutionProvider 기준 검증

한계

• 단일 라벨만 출력합니다. "사진 보여주고 기사도 찾아줘" 같은 멀티 의도 발화는 학습 데이터에서 의도적으로 배제했으며 지원하지 않습니다.
• 전량 합성 데이터로 학습했습니다. 실제 사용자 로그 기반 평가는 이루어지지 않았습니다.
• confidence를 신뢰도 필터로 쓰기 어렵습니다. 오분류 시에도 높은 확신도(0.89+)를 보이는 경향이 있어, confidence 임계값 기반 폴백은 권장하지 않습니다.
• 입력은 한국어 단문(대화 발화) 기준입니다. 긴 문서나 타 언어 입력은 검증되지 않았습니다.

파일

라이선스

Apache-2.0 (베이스 모델 KoELECTRA와 동일)