Huggingface datasets を使って オリジナルデータでNER

BERTも使うし、どうせ後でコードが散らかるよりはということで、
huggingface の datasetsライブラリを使ってみる

とりあえず辞書から

from datasets import Dataset, DatasetDict
tokens = [["I", "love", "you", "."], ["I", "hate", "you", "."]]
ner_tags = [[0, 1, 0, 0], [0, 1, 0, 0]] 
# [["O", "B-Aspect", "O", "O"], ["O", "B-Aspect", "O", "O"]]
train = Dataset.from_dict({"tokens": tokens, "ner_tags": ner_tags})
DatasetDict({"train":train})

一応これでDataset自体は整形できる

型情報追加

上のままだと、型情報を見てくれないので、型情報を追加する。

こうしておくと、transformers側で logitじゃなくてlabelを予測してくれるなどのメリットがあるほか、datasets に元から登録されている NER のデータセットのように共通して扱えるメリットがある。

手順は次のとおり

• datasets.Features を利用して、DatasetInfoを作る
• Datasetにinfoを渡す

import datasets
ner_label_names = ["O", "B-Aspect", "I-Aspect"]
info = datasets.DatasetInfo(
                features=datasets.Features(
                    {
                        "tokens": datasets.Sequence(datasets.Value("string")),
                        "ner_tags": datasets.Sequence(
                            datasets.ClassLabel(names=ner_label_names)
                        ),
                    }
                )
)
train = Dataset.from_dict(
                {"tokens": tokens, "ner_tags": ner_tags}, info=info
            )

参考資料

• https://huggingface.co/docs/datasets/loading_datasets.html
• https://huggingface.co/docs/datasets/add_dataset.html

本当はDatasetBuilderっていうのを使うらしいけど、使い方がよくわからない...