gihakkk/tinybert_model

사용예시

import onnxruntime as ort
import numpy as np
from transformers import MobileBertTokenizer

# 모델 및 토크나이저 경로 설정
model_path = r'C:\NEW_tinybert\AI\tinybert_model.onnx'  # ONNX 모델 경로
tokenizer_path = r'C:\NEW_distilbert\AI'  # 로컬 토크나이저 경로

# ONNX 모델 세션 초기화
ort_session = ort.InferenceSession(model_path)

# MobileBERT 토크나이저 로드
tokenizer = MobileBertTokenizer.from_pretrained(tokenizer_path)

# 텍스트 분류 함수
def test_model(text):
    """
    입력된 텍스트를 ONNX 모델을 사용해 분류하는 함수
    Args:
        text (str): 분석할 텍스트
    Returns:
        str: 예측 결과 메시지
    """
    # 입력 텍스트를 토큰화 및 ONNX 모델 입력 형식으로 변환
    inputs = tokenizer(
        text,
        padding="max_length",  # 입력 길이를 128로 고정
        truncation=True,       # 긴 텍스트는 잘라냄
        max_length=128,        # 최대 토큰 길이
        return_tensors="np"    # NumPy 배열 형식으로 반환
    )
    
    # NumPy 배열을 int64로 변환
    input_ids = inputs["input_ids"].astype(np.int64)
    attention_mask = inputs["attention_mask"].astype(np.int64)
    
    # ONNX 모델 입력 준비
    ort_inputs = {
        "input_ids": input_ids,
        "attention_mask": attention_mask
    }
    
    # ONNX 모델 추론 실행
    outputs = ort_session.run(None, ort_inputs)
    logits = outputs[0]  # 모델 출력 (로짓 값)

    # 로짓 값을 확률로 변환 및 클래스 예측
    predicted_class = np.argmax(logits, axis=1).item()

    # 결과 반환
    return "로맨스 스캠일 가능성 있음" if predicted_class == 1 else "로맨스 스캠이 아님"

# 테스트할 대화 내용
test_texts = [
    "너 엄마 없냐?",
    "저는 금융 상품을 소개하는 사람입니다. 투자하면 이익이 큽니다.",
    "내 보지가 달아올랐어",
    "내 가슴 만질래??"
]

# 각 테스트 텍스트에 대해 결과 출력
for text in test_texts:
    result = test_model(text)
    print(f"입력: {text} => 결과: {result}")