TransformerでTokenizeしたSpanを入力文から特定

Hugging faceで文をTokenizeすると元の文のSpan情報が失われてしまう。これにより、NERでラベル付けしたトークンの位置をサブワード分割された文の中で見失ってしまうことがある。この困りごとを解決してくれるのが、return_offsets_mappingオプションである。

次のコードで、このオプションの使い方と、サブワード分割されたトークンと、元のトークンの対応関係を出力する方法を紹介。

from transformers import XLMRobertaTokenizerFast
tokenizer = XLMRobertaTokenizerFast.from_pretrained('xlm-roberta-large')
text = "It builds on BERT and modifies key hyperparameters"
tokenized = tokenizer(text, return_offsets_mapping=True)
tokenized_ = tokenizer.tokenize(text)
print('Text:', text)
print('Tokenized:', tokenized_)
print('Tokens:', tokenizer.convert_ids_to_tokens(tokenized.input_ids))
print('Mapped to:', [text[start:end] for start, end in tokenized.offset_mapping])

from transformers import AutoModel, AutoTokenizer, AutoModelForPreTraining
tokenizer = AutoTokenizer.from_pretrained("allenai/scibert_scivocab_uncased")
text = "It builds on BERT and modifies key hyperparameters"
tokenized = tokenizer(text, return_offsets_mapping=True)
tokenized_ = tokenizer.tokenize(text)
print('Text:', text)
print('Tokenized:', tokenized_)
print('Tokens:', tokenizer.convert_ids_to_tokens(tokenized.input_ids))
print('Mapped to:', [text[start:end] for start, end in tokenized.offset_mapping])

以下出力結果。

Text: It builds on BERT and modifies key hyperparameters
Tokenized: ['▁It', '▁build', 's', '▁on', '▁BER', 'T', '▁and', '▁modifi', 'es', '▁key', '▁hyper', 'para', 'meter', 's']
Tokens: ['<s>', '▁It', '▁build', 's', '▁on', '▁BER', 'T', '▁and', '▁modifi', 'es', '▁key', '▁hyper', 'para', 'meter', 's', '</s>']
Mapped to: ['', 'It', 'build', 's', 'on', 'BER', 'T', 'and', 'modifi', 'es', 'key', 'hyper', 'para', 'meter', 's', '']

Text: It builds on BERT and modifies key hyperparameters
Tokenized: ['it', 'builds', 'on', 'bert', 'and', 'modifies', 'key', 'hyper', '##parameters']
Tokens: ['[CLS]', 'it', 'builds', 'on', 'bert', 'and', 'modifies', 'key', 'hyper', '##parameters', '[SEP]']
Mapped to: ['', 'It', 'builds', 'on', 'BERT', 'and', 'modifies', 'key', 'hyper', 'parameters', '']
Pocket