토큰 분류

토큰 분류 (Token classification)

토큰 분류는 문장 안의 개별 토큰에 라벨을 부여하는 작업이에요. 가장 흔한 토큰 분류 작업 중 하나는 개체명 인식(NER, Named Entity Recognition)이에요. NER은 문장 안의 각 개체(entity)에 대해 사람, 위치, 조직 같은 라벨을 찾아내요.

출처: 문서

본문

이 가이드에서는 다음을 배워요:

  1. WNUT 17 데이터셋에서 DistilBERT를 파인튜닝해서 새로운 개체를 감지하기.
  2. 파인튜닝한 모델을 추론(inference)에 사용하기.

이 작업과 호환되는 모든 아키텍처와 체크포인트를 보려면 task-page를 확인하는 걸 추천해요.

시작하기 전에 필요한 라이브러리를 모두 설치했는지 확인해요:

pip install transformers datasets evaluate seqeval

Hugging Face 계정에 로그인해서 모델을 커뮤니티에 업로드하고 공유하는 걸 권장해요. 로그인하라는 메시지가 나오면 토큰을 입력해서 로그인하세요:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

WNUT 17 데이터셋 로드

먼저 🤗 Datasets 라이브러리에서 WNUT 17 데이터셋을 로드해요:

>>> from datasets import load_dataset

>>> wnut = load_dataset("wnut_17")

그다음 예시 하나를 살펴봐요:

>>> wnut["train"][0]
{'id': '0',
 'ner_tags': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 8, 8, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0],
 'tokens': ['@paulwalk', 'It', "'s", 'the', 'view', 'from', 'where', 'I', "'m", 'living', 'for', 'two', 'weeks', '.', 'Empire', 'State', 'Building', '=', 'ESB', '.', 'Pretty', 'bad', 'storm', 'here', 'last', 'evening', '.']}

ner_tags의 각 숫자는 하나의 개체를 나타내요. 숫자를 라벨 이름으로 변환하면 개체가 무엇인지 알 수 있어요:

>>> label_list = wnut["train"].features[f"ner_tags"].feature.names
>>> label_list
[
    "O",
    "B-corporation",
    "I-corporation",
    "B-creative-work",
    "I-creative-work",
    "B-group",
    "I-group",
    "B-location",
    "I-location",
    "B-person",
    "I-person",
    "B-product",
    "I-product",
]

각 ner_tag 앞에 붙는 알파벳은 개체의 토큰 위치를 나타내요:

  • B-는 개체의 시작을 나타내요.
  • I-는 토큰이 같은 개체 안에 포함되어 있음을 나타내요 (예: State 토큰은 Empire State Building 같은 개체의 일부).
  • 0은 토큰이 어떤 개체에도 해당하지 않는다는 걸 나타내요.

전처리

다음 단계는 tokens 필드를 전처리할 DistilBERT 토크나이저를 로드하는 거예요:

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")

위의 tokens 필드 예시에서 봤듯이 입력이 이미 토크나이즈된 것처럼 보여요. 하지만 실제로는 아직 토크나이즈되지 않았고, 단어를 서브워드로 토크나이즈하려면 is_split_into_words=True를 설정해야 해요. 예를 들어:

>>> example = wnut["train"][0]
>>> tokenized_input = tokenizer(example["tokens"], is_split_into_words=True)
>>> tokens = tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])
>>> tokens
['[CLS]', '@', 'paul', '##walk', 'it', "'", 's', 'the', 'view', 'from', 'where', 'i', "'", 'm', 'living', 'for', 'two', 'weeks', '.', 'empire', 'state', 'building', '=', 'es', '##b', '.', 'pretty', 'bad', 'storm', 'here', 'last', 'evening', '.', '[SEP]']

그런데 이렇게 하면 [CLS]와 [SEP] 같은 특수 토큰이 추가되고, 서브워드 토크나이제이션 때문에 입력과 라벨 사이에 불일치가 생겨요. 하나의 라벨에 해당하는 한 단어가 두 개의 서브워드로 나뉠 수 있기 때문이에요. 토큰과 라벨을 다시 정렬하려면 다음을 해야 해요:

  1. word_ids 메서드로 모든 토큰을 해당하는 단어에 매핑해요.
  2. 특수 토큰 [CLS]와 [SEP]에 라벨 -100을 할당해서 PyTorch 손실 함수가 무시하게 해요 (CrossEntropyLoss 참고).
  3. 주어진 단어의 첫 번째 토큰에만 라벨을 붙이고, 같은 단어의 나머지 서브토큰에는 -100을 할당해요.

토큰과 라벨을 다시 정렬하는 함수를 만들고, 시퀀스가 DistilBERT의 최대 입력 길이를 넘지 않도록 자르는 방법은 다음과 같아요:

>>> def tokenize_and_align_labels(examples):
...     tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)

...     labels = []
...     for i, label in enumerate(examples[f"ner_tags"]):
...         word_ids = tokenized_inputs.word_ids(batch_index=i)  # Map tokens to their respective word.
...         previous_word_idx = None
...         label_ids = []
...         for word_idx in word_ids:  # Set the special tokens to -100.
...             if word_idx is None:
...                 label_ids.append(-100)
...             elif word_idx != previous_word_idx:  # Only label the first token of a given word.
...                 label_ids.append(label[word_idx])
...             else:
...                 label_ids.append(-100)
...             previous_word_idx = word_idx
...         labels.append(label_ids)

...     tokenized_inputs["labels"] = labels
...     return tokenized_inputs

전처리 함수를 전체 데이터셋에 적용하려면 🤗 Datasets의 map 함수를 사용해요. batched=True를 설정하면 데이터셋의 여러 요소를 한 번에 처리할 수 있어서 map 함수를 빠르게 할 수 있어요:

>>> tokenized_wnut = wnut.map(tokenize_and_align_labels, batched=True)

이제 DataCollatorWithPadding을 사용해 예시 배치를 만들어요. 콜레이션(collation) 과정에서 데이터셋 전체를 최대 길이로 패딩하는 것보다 배치 안의 문장들을 가장 긴 길이에 맞춰 동적으로 패딩하는 편이 더 효율적이에요.

>>> from transformers import DataCollatorForTokenClassification

>>> data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

평가

학습 중에 메트릭을 포함하면 모델의 성능을 평가하는 데 종종 도움이 돼요. 🤗 Evaluate 라이브러리로 평가 방법을 빠르게 로드할 수 있어요. 이 작업에서는 seqeval 프레임워크를 로드해요 (메트릭을 로드하고 계산하는 방법은 🤗 Evaluate quick tour를 참고하세요). seqeval은 실제로 precision, recall, F1, accuracy 등 여러 점수를 만들어 내요.

>>> import evaluate

>>> seqeval = evaluate.load("seqeval")

먼저 NER 라벨을 가져오고, 그다음 실제 예측값과 실제 라벨을 compute에 넘겨 점수를 계산하는 함수를 만들어요:

>>> import numpy as np

>>> labels = [label_list[i] for i in example[f"ner_tags"]]

>>> def compute_metrics(p):
...     predictions, labels = p
...     predictions = np.argmax(predictions, axis=2)

...     true_predictions = [
...         [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
...         for prediction, label in zip(predictions, labels)
...     ]
...     true_labels = [
...         [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
...         for prediction, label in zip(predictions, labels)
...     ]

...     results = seqeval.compute(predictions=true_predictions, references=true_labels)
...     return {
...         "precision": results["overall_precision"],
...         "recall": results["overall_recall"],
...         "f1": results["overall_f1"],
...         "accuracy": results["overall_accuracy"],
...     }

이제 compute_metrics 함수를 쓸 준비가 됐어요. 학습을 설정할 때 다시 사용하게 될 거예요.

학습

모델 학습을 시작하기 전에 기대되는 id와 라벨 사이의 매핑을 id2label, label2id로 만들어요:

>>> id2label = {
...     0: "O",
...     1: "B-corporation",
...     2: "I-corporation",
...     3: "B-creative-work",
...     4: "I-creative-work",
...     5: "B-group",
...     6: "I-group",
...     7: "B-location",
...     8: "I-location",
...     9: "B-person",
...     10: "I-person",
...     11: "B-product",
...     12: "I-product",
... }
>>> label2id = {
...     "O": 0,
...     "B-corporation": 1,
...     "I-corporation": 2,
...     "B-creative-work": 3,
...     "I-creative-work": 4,
...     "B-group": 5,
...     "I-group": 6,
...     "B-location": 7,
...     "I-location": 8,
...     "B-person": 9,
...     "I-person": 10,
...     "B-product": 11,
...     "I-product": 12,
... }

Trainer로 모델을 파인튜닝하는 방법에 익숙하지 않다면 여기의 기본 튜토리얼을 확인해 보세요!

이제 모델 학습을 시작할 준비가 됐어요! DistilBERT를 AutoModelForTokenClassification로 로드하고, 기대되는 라벨 개수와 라벨 매핑을 함께 전달해요:

>>> from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer

>>> model = AutoModelForTokenClassification.from_pretrained(
...     "distilbert/distilbert-base-uncased", num_labels=13, id2label=id2label, label2id=label2id
... )

이 시점에서 남은 단계는 세 가지뿐이에요:

  1. TrainingArguments에서 학습 하이퍼파라미터를 정의해요. 유일하게 필수인 파라미터는 모델을 저장할 위치를 지정하는 output_dir이에요. push_to_hub=True로 설정하면 모델을 Hub에 푸시할 수 있어요 (모델을 업로드하려면 Hugging Face에 로그인해야 해요). 각 에폭이 끝날 때마다 Trainer가 seqeval 점수를 평가하고 학습 체크포인트를 저장해요.
  2. 학습 인수를 Trainer에 모델, 데이터셋, 토크나이저, 데이터 콜레이터, compute_metrics 함수와 함께 전달해요.
  3. train()을 호출해서 모델을 파인튜닝해요.
>>> training_args = TrainingArguments(
...     output_dir="my_awesome_wnut_model",
...     learning_rate=2e-5,
...     per_device_train_batch_size=16,
...     per_device_eval_batch_size=16,
...     num_train_epochs=2,
...     weight_decay=0.01,
...     eval_strategy="epoch",
...     save_strategy="epoch",
...     load_best_model_at_end=True,
...     push_to_hub=True,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=tokenized_wnut["train"],
...     eval_dataset=tokenized_wnut["test"],
...     processing_class=tokenizer,
...     data_collator=data_collator,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

학습이 완료되면 push_to_hub() 메서드로 모델을 Hub에 공유해서 모두가 쓸 수 있게 해요:

>>> trainer.push_to_hub()

토큰 분류용으로 모델을 파인튜닝하는 더 심층적인 예시는 해당 PyTorch notebook을 참고하세요.

추론 (Inference)

좋아요, 이제 모델을 파인튜닝했으니 추론에 활용할 수 있어요!

추론을 실행하고 싶은 텍스트를 가져와요:

>>> text = "The Golden State Warriors are an American professional basketball team based in San Francisco."

파인튜닝한 모델을 추론에 사용하는 가장 간단한 방법은 pipeline()에서 쓰는 거예요. 모델로 NER용 pipeline을 만들고 텍스트를 전달해요:

>>> from transformers import pipeline

>>> classifier = pipeline("ner", model="stevhliu/my_awesome_wnut_model")
>>> classifier(text)
[{'entity': 'B-location',
  'score': 0.42658573,
  'index': 2,
  'word': 'golden',
  'start': 4,
  'end': 10},
 {'entity': 'I-location',
  'score': 0.35856336,
  'index': 3,
  'word': 'state',
  'start': 11,
  'end': 16},
 {'entity': 'B-group',
  'score': 0.3064001,
  'index': 4,
  'word': 'warriors',
  'start': 17,
  'end': 25},
 {'entity': 'B-location',
  'score': 0.65523505,
  'index': 13,
  'word': 'san',
  'start': 80,
  'end': 83},
 {'entity': 'B-location',
  'score': 0.4668663,
  'index': 14,
  'word': 'francisco',
  'start': 84,
  'end': 93}]

원한다면 pipeline의 결과를 직접 재현할 수도 있어요:

텍스트를 토크나이즈하고 PyTorch 텐서로 반환해요:

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("stevhliu/my_awesome_wnut_model")
>>> inputs = tokenizer(text, return_tensors="pt")

입력을 모델에 전달하고 logits를 반환해요:

>>> from transformers import AutoModelForTokenClassification

>>> model = AutoModelForTokenClassification.from_pretrained("stevhliu/my_awesome_wnut_model")
>>> with torch.no_grad():
...     logits = model(**inputs).logits

가장 높은 확률을 가진 클래스를 가져오고, 모델의 id2label 매핑을 사용해 텍스트 라벨로 변환해요:

>>> predictions = torch.argmax(logits, dim=2)
>>> predicted_token_class = [model.config.id2label[t.item()] for t in predictions[0]]
>>> predicted_token_class
['O',
 'O',
 'B-location',
 'I-location',
 'B-group',
 'O',
 'O',
 'O',
 'O',
 'O',
 'O',
 'O',
 'O',
 'B-location',
 'B-location',
 'O',
 'O']