yasuokaの日記： Transformersにおける日本語トークナイザBertJapaneseTokenizerFastの試作

昨年12月26日の日記で、私(安岡孝一)は以下のように書いた。

一方で、fugashi(というかBertJapaneseTokenizer)をトークナイザに使うと、オプションにreturn_offsets_mapping=Trueが効かないため、11月26日の日記で書いたような手法が適用できない。うーん、国語研短単位をサポートするようなPreTrainedTokenizerFastを、何とかデッチ上げるしかないかな。

元旦の宿題に、PreTrainedTokenizerFastとまではいかなかったものの、何とかBertJapaneseTokenizerFastっぽいクラスを、pytokenizationsでデッチ上げてみた。Google Colaboratoryで動かしてみよう。

!pip install transformers pytokenizations fugashi unidic_lite
from transformers import BertJapaneseTokenizer
class BertJapaneseTokenizerFast(BertJapaneseTokenizer):
  def __call__(self,text,text_pair=None,return_offsets_mapping=False,**kwargs):
    v=super().__call__(text=text,text_pair=text_pair,return_offsets_mapping=False,**kwargs)
    if return_offsets_mapping:
      import tokenizations
      if type(text)==str:
        a=v["input_ids"]
        b=text+text_pair if text_pair else text
        a2b,b2a=tokenizations.get_alignments(self.convert_ids_to_tokens(a),b)
        w=[]
        for i,t in enumerate(a2b):
          if t==[]:
            s=(0,0)
            if a[i]==self.unk_token_id:
              if i>0 and i<len(a2b)-1 and a2b[i-1]!=[] and a2b[i+1]!=[]:
                s=(a2b[i-1][-1]+1,a2b[i+1][0])
          elif t[-1]<len(text):
            s=(t[0],t[-1]+1)
          else:
            s=(t[0]-len(text),t[-1]-len(text)+1)
          w.append(s)
        v["offset_mapping"]=w
      else:
        w=[]
        for a,b in zip(v["input_ids"],text+text_pair if text_pair else text):
          a2b,b2a=tokenizations.get_alignments(self.convert_ids_to_tokens(a),b)
          x=[]
          for i,t in enumerate(a2b):
            if t==[]:
              s=(0,0)
              if a[i]==self.unk_token_id:
                if i>0 and i<len(a2b)-1 and a2b[i-1]!=[] and a2b[i+1]!=[]:
                  s=(a2b[i-1][-1]+1,a2b[i+1][0])
            else:
              s=(t[0],t[-1]+1)
            x.append(s)
          w.append(list(x))
        v["offset_mapping"]=w
    return v

tokenizer=BertJapaneseTokenizerFast.from_pretrained("cl-tohoku/bert-base-japanese-v2")
d=tokenizer("国境の長いトンネルを抜けると雪国であった。","夜の底が白くなった。",return_offsets_mapping=True)
print(tokenizer.convert_ids_to_tokens(d["input_ids"]),d,sep="\n")

cl-tohoku/bert-base-japanese-v2のトークナイザを拡張して、「国境の長いトンネルを抜けると雪国であった。」と「夜の底が白くなった。」をトークナイズしてみたところ、私の手元では以下の結果になった。

['[CLS]', '国境', 'の', '長い', 'トンネル', 'を', '抜ける', 'と', '雪', '##国', 'で', 'あっ', 'た', '。', '[SEP]', '夜', 'の', '底', 'が', '白く', 'なっ', 'た', '。', '[SEP]']
{'input_ids': [2, 14852, 896, 12957, 13762, 932, 27762, 890, 5593, 6368, 889, 11170, 881, 829, 3, 1844, 896, 2197, 862, 27262, 11147, 881, 829, 3], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 'offset_mapping': [(0, 0), (0, 2), (2, 3), (3, 5), (5, 9), (9, 10), (10, 13), (13, 14), (14, 15), (15, 16), (16, 17), (17, 19), (19, 20), (20, 21), (0, 0), (0, 1), (1, 2), (2, 3), (3, 4), (4, 6), (6, 8), (8, 9), (9, 10), (0, 0)]}

offset_mappingがちゃんと出ていて、かなりうれしい。ただ、この手法だと[UNK]に対する処理が甘く、[UNK]の前後に[CLS]や[SEP]や[UNK]が来た場合にうまくいかない。さて、このあたり、どうしたらいいかな。