검색
검색
공개 노트 검색
회원가입로그인

브라우저에서 RoBERTa 실행해보기

RoBERTa를 사용하면 텍스트 분류, 검색, 특징 탐지과 같은 기능을 구현할 수 있습니다. 기술의 발전으로 인해 가벼운 언어모델들은 브라우저에서 직접 실행할 수 있게 됐습니다.

이렇게 인공지능 모델을 사용자 디바이스에서 직접 실행하는 기술을 온디바이스 AI(On-device AI)라고 합니다.

일상을 바꿀 내 손 안의 AI, 온디바이스 AI – 제일기획 매거진

Transformers.js

자바스크립트 런타임에서 Transformers 언어모델을 실행할 수 있게 하는 라이브러리입니다. 특히 브라우저에서도 실행이 가능합니다. 이 라이브러리를 사용해서 구현할 예정입니다. Transformers.js 라이브러리는 ONNX 포맷으로 저장된 모델을 로드할 수 있습니다.

하지만 대부분의 언어 모델은 pth(PyTorch), h5(TensorFlow), safetensors(SafeTensors)로 배포됩니다. 그렇기에 ONNX로 변환을 해야 Transformers.js에서 사용할 수 있습니다.

Transformers 모델을 브라우저가 이해할 수 있도록 PyTorch ONNX로 변환하기

시범으로 leewaay/klue-roberta-base-klueNLI-klueSTS-MSL128 모델을 변환해보겠습니다. Huggingface의 Optimum이라는 라이브러리를 사용하면 손쉽게 변환할 수 있습니다.

먼저 필요한 라이브러리를 설치해주겠습니다.

pip install optimum onnxruntime sentence_transformers onnx

언어모델의 가중치를 ONNX로 변환해보겠습니다

from optimum.onnxruntime import ORTModelForSequenceClassification, ORTModelForFeatureExtraction
from transformers import AutoTokenizer

model_checkpoint = "leewaay/klue-roberta-base-klueNLI-klueSTS-MSL128"
save_directory = "tmp/onnx/"

ort_model = ORTModelForFeatureExtraction.from_pretrained(model_checkpoint, export=True)
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

ort_model.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)

1분 내외로 ./tmp/onnx 디렉토리에 ONNX 모델 파일이 저장됩니다. 다음과 같은 결과물을 확인할 수 있습니다.

브라우저에서 RoBERTa 실행해보기 image 1

  • tmp/onnx

    • config.json

    • model.onnx

    • special_tokens_map

    • tokenizer_config.json

    • tokenizer.json

    • vocab.txt

(선택) 양자화로 모델 크기 줄이기

400MB는 앱의 일부로 배포하기에는 부담스러운 크기입니다. 모델의 크기도 줄이고 추론 속도도 높히려면 양자화를 수행할 수 있습니다. 다음 코드로 모델을 양자화할 수 있습니다.

같이 보기: 딥러닝 모델 쉽게 양자화하기

양자화는 딥러닝 모델을 경량화하는 방법중 하나입니다. Huggingface의 transformers 라이브러리가 제공하는 툴킷을 사용하면 Huggingface에 업로드된 모델을 간편하게 양자화할 수 있습니다...

from optimum.onnxruntime import ORTQuantizer, AutoQuantizationConfig

quantizer = ORTQuantizer.from_pretrained(ort_model)
dqconfig = AutoQuantizationConfig.avx512_vnni(
  is_static = False,
  per_channel = False
)

model_quantized_path = quantizer.quantize(
    save_dir = save_directory,
    quantization_config = dqconfig,
)

위 코드를 실행하면 `save_directory`에 `model_quantized.onnx`라는 파일이 추가로 생성됩니다. 양자화 이후에 크기가 110MB 가량으로 줄었습니다.

Transformers.js에서 ONNX 모델을 실행하기

transformers.js를 설치해주겠습니다.

pnpm install @xenova/transformers

다음과 같이 모델을 로드할 수 있습니다.

import { AutoTokenizer, AutoModel } from '@xenova/transformers'

const modelName = 'model-directory'

const tokenizer = await AutoTokenizer.from_pretrained(modelName)
const model = await AutoModel.from_pretrained(modelName, {
    quantized: false, // 양자화를 수행했다면 생략할 수 있음
})

문장 임베딩 계산하기

Sentence Transformers 모델은 mean pooling으로 문장 임베딩을 계산할 수 있다.

import { mean_pooling } from '@xenova/transformers'

async function getSentenceEmbedding(text) {
  const inputs = await tokenizer(text)
  const result = await model(inputs)
  
  const attentionMask = inputs.attention_mask
  const pooled = mean_pooling(result.last_hidden_state, attentionMask)

  return pooled.data
}

문장 간 유사도 계산하기

import { cos_sim } from '@xenova/transformers'

const s1 = await getSentenceEmbedding("난 대학시절 묵찌빠를 전공했단 사실")
const s2 = await getSentenceEmbedding("난 묵찌빠로 유학까지 다녀왔단 사실")
const s3 = await getSentenceEmbedding('네 놈을 이겨 가문의 이름 높이리')

console.log(cos_sim(s1, s2)) // 0.6303152271221659
console.log(cos_sim(s1, s3)) // 0.13926317908806657
공유하기
카카오로 공유하기
페이스북 공유하기
트위터로 공유하기
url 복사하기
조회수 : 511
heart
T
페이지 기반 대답
AI Chat