DistilBERT 실전 사용 — 파이프라인과 파인튜닝

DistilBERT 실전 사용

DistilBERT는 원시 모델로는 마스크드 언어 모델링이나 차기 문장 예측에 쓸 수 있지만, 실제로는 다운스트림 태스크에 파인튜닝해서 쓰는 게 주 목적이에요. 시퀀스 분류, 토큰 분류, 질의응답처럼 문장 전체를 보고 판단하는 태스크에 특히 잘 어울려요. 텍스트 생성 용도라면 GPT2 같은 모델을 찾는 편이 좋아요.

출처: https://huggingface.co/distilbert/distilbert-base-uncased

파이프라인으로 분류하기

from transformers import pipeline

classifier = pipeline(
    task="text-classification",
    model="distilbert-base-uncased-finetuned-sst-2-english",
    device=0
)

result = classifier("I love using Hugging Face Transformers!")
print(result)
# Output: [{'label': 'POSITIVE', 'score': 0.9998}]

마스크드 언어 모델로 쓰기

from transformers import pipeline
unmasker = pipeline('fill-mask', model='distilbert-base-uncased')
unmasker("Hello I'm a [MASK] model.")

PyTorch / TensorFlow로 임베딩 추출

from transformers import DistilBertTokenizer, DistilBertModel
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertModel.from_pretrained("distilbert-base-uncased")
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)

평가 결과 (GLUE)

파인튜닝했을 때 GLUE 테스트 성적이에요.

Task MNLI QQP QNLI SST-2 CoLA STS-B MRPC RTE
82.2 88.5 89.2 91.3 51.3 85.8 87.5 59.9

더 알아보기