#はじめに
huggingface/transformersのpipelineは数行のコードで推論を行えますが、出力が少々不自由です。
本記事ではpipelineの一部を変更し、出力にtokenized textとattention weightを加えました。
具体的には使用したいpipeline(本記事ではTextClassificationPipeline
)を継承し、クラスメソッドのforword
とpostprocess
に数行のコードを加えます。これにより、自作のコードを最小限にしつつ欲しい出力を得られます。
本記事ではBERTのtext classificationを例にして、下記の構成で説明します。
- tokenized textの取得方法
- attention weightの取得方法
- pipelineでのtokenized textとattention weightの出力方法 (本題)
コードだけを見たい方はこちらまで飛ばしてお読みください。
huggingface/transformersの基本的な使用方法は説明しませんので、公式ドキュメントをご確認ください。
また筆者が未熟なため、記事内容の誤りや記事内容より優れた方法があるかと思います。そのような箇所を見つけた際にはコメントでご指摘いただけますと幸いです。
#tokenized textの取得方法
まずは下記の通りtokenizerでtextを分割します。
from transformers import AutoTokenizer
model_name = '学習済みモデルのpath'
tokenizer = AutoTokenizer.from_pretrained(model_name)
#「吾輩は猫である」を例文に使用します
text = '吾輩は猫である。名前はまだ無い。'
model_inputs = tokenizer(text, return_tensors='pt')
{
'input_ids': tensor([[ 2, 1583, 5159, 897, 3574, 889, 20656, 829, 1564, 1402, 897, 6466, 6626, 3464, 854, 829, 3]]),
'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]),
'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
}
input_ids
からtokenを復元すれば、tokenized textが取得可能だとわかります。
huggingface/transformersには、このためのメソッド(convert_ids_to_tokens
)があらかじめ用意されています。
tokenized_text = tokenizer.convert_ids_to_tokens(model_inputs['input_ids'][0])
['[CLS]', '吾', '輩', 'は', '猫', 'で', '##ある', '。', '名', '前', 'は', '##ま', '##だ', '無', 'い', '。', '[SEP]']
以上で、tokenized textの取得方法がわかりました。
attention weightの取得方法
モデル設定でoutput_attentions=True
を渡すことで、出力結果にattention weightが含まれます。
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(model_name, output_attentions=True)
model_outputs = model(model_inputs['input_ids'])
# attentionsは省略して記載しているため、実際の出力とは異なります
SequenceClassifierOutput(
loss=None,
logits=tensor([[0.0077, 0.1880]], grad_fn=<AddmmBackward0>),
hidden_states=None,
attentions=(tensor([[[[0.1091, 0.0460, 0.0619, ..., 0.0487, 0.0344, 0.0468],[0.0431, 0.0089, 0.1076, ..., 0.1150, 0.0315, 0.0183],[0.1115, 0.0604, 0.0373, ..., 0.0246, 0.0797, 0.1302], ...,[4.1084e-03, 9.1141e-04, 1.6878e-04, ..., 7.1091e-03, 1.2396e-02, 9.3627e-01], [1.8521e-02, 3.7775e-03, 5.7404e-03, ..., 6.6644e-03, 1.0438e-02, 9.2752e-01], [3.5021e-02, 1.7682e-02, 3.3329e-02, ..., 1.2351e-02, 1.4987e-02, 7.6510e-01]]]], grad_fn=<SoftmaxBackward0>))
)
つまり、model_outputs['attentions']
でattention weightが取得できます。
#pipelineでのtokenized textとattention weightの出力方法
本題に入ります。
##pipelineの出力結果
まずはpipelineの出力を確認します。
from transformers import TextClassificationPipeline
classifier = TextClassificationPipeline(model=model, tokenizer=tokenizer)
outputs = classifier(text)
[{'label': 'LABEL_1', 'score': 0.5449662804603577}]
このようにcode3
でoutput_attentions=True
を渡していても、pipelineの出力結果にはattention weightがありません。これはクラスメソッドのpostprocess
によるものです。
##postprocessの説明およびattention weightの出力方法
postprocess
は、生の出力(model_outputs
)を受け取り、見やすい結果(outputs
)を返すメソッドです。しかし、返り値にattention weightがないため、result4
のようになります。
つまり、postprocess
の返り値にattention weightを加えることで、attention weightを出力できます。
##tokenized textの出力方法
postprocess
の返り値を変更すれば、出力結果を自由にできることがわかりました。しかし、code2
のようにtokenized textを得るにはmodel_inputs
が必要です。また、model_inputs
から取得したtokenized textをpostprocess
に渡さなければなりません。そのため、model_inputs
を受け取り、model_outputs
を返すクラスメソッドforward
を変更します。
forward
が受け取ったmodel_inputs
からcode2
でtokenized textを取得し、model_outputs
に加えてpostprocess
に引き渡します。
後は、attention weightと同様にpostprocess
を変更すれば、tokenized textも出力できます。
##変更を加えたクラスの実装
以上を踏まえて、forward
とpostprocess
に変更を加えます。
huggingface/transformersを直接修正しても良いですが、元のライブラリに変更は加えたくないため、
TextClassificationPipeline
を継承しTextClassificationPipelineAddOutputs
を実装します。
class TextClassificationPipelineAddOutputs(TextClassificationPipeline):
def forward(self, model_inputs, **forward_params):
# 既存の処理を行います
model_outputs = super().forward(model_inputs, **forward_params)
# code2と同様です
tokenized_text = self.tokenizer.convert_ids_to_tokens(model_inputs['input_ids'][0])
# model_outputs(辞書型)にtokenized textを加えます
model_outputs['tokenized_text'] = tokenized_text
return model_outputs
def postprocess(self, model_outputs, function_to_apply=None, return_all_scores=False):
# 既存の処理を行います
processed = super().postprocess(model_outputs, function_to_apply, return_all_scores)
# 見やすさのために変数を用意しています
tokenized_text = model_outputs['tokenized_text']
attentions = model_outputs['attentions']
# outputsにtokenized textとattention weightを追加
if return_all_scores:
return [{**proc, **{'tokenized_text':tokenized_text, 'attentions':attentions}} for proc in processed]
else:
return {**processed, **{'tokenized_text':tokenized_text, 'attentions':attentions}}
実装したクラスの出力を確認します。
# code3と同様に、'model'に'output_attentions=True'を渡してください
classifier = TextClassificationPipelineAddOutputs(model=model, tokenizer=tokenizer)
outputs = classifier(text)
# attentionsは省略して記載しているため、実際の出力とは異なります
[{
'label': 'LABEL_1',
'score': 0.5449662804603577,
'tokenized_text': ['[CLS]', '吾', '輩', 'は', '猫', 'で', '##ある', '。', '名', '前', 'は', '##ま', '##だ', '無', 'い', '。', '[SEP]'],
'attentions': (tensor([[[[0.1091, 0.0460, 0.0619, ..., 0.0487, 0.0344, 0.0468],[0.0431, 0.0089, 0.1076, ..., 0.1150, 0.0315, 0.0183],[0.1115, 0.0604, 0.0373, ..., 0.0246, 0.0797, 0.1302], ..., [0.0074, 0.1923, 0.2223, ..., 0.0083, 0.0186, 0.0239], [0.0107, 0.0105, 0.0127, ..., 0.0151, 0.1036, 0.6139], [0.0220, 0.0112, 0.0168, ..., 0.0362, 0.3030, 0.1718]]]]))}]
tokenized textとattention weightが出力できました。
#おわりに
本記事がお役に立てれば幸いです。
繰り返しになりますが、記事内容の誤りや記事内容より優れた方法があれば、コメントでご指摘いただけますとありがたく存じます。