日本語のテキストコーパスから辞書を作るライブラリを作りたい

DeepLearningを使って自然言語処理をする際に、今まで毎回毎回テキストから辞書を作るコードをゼロから書いてきました。

毎回面倒に感じていたので、今回テキストを用意したら一瞬で辞書を作ってくれるプログラムを書いてみました。

環境
作ったもの
使い方
- vocab
- word2id
- id2word
- encode
- decode
ソースコード
ゼロから作るDeepLearning
続き

環境

※MeCabをWindowsにインストールするには

https://www.pytry3g.com/entry/2018/04/24/143934#Windows

作ったもの

このプログラムはPyTorch-NLPのプログラムをもろパクリ参考にして作りました。テキストコーパスから辞書を作り、テキストを与えると、単語IDが入ったリストを返してくれます。

>>> corpus = ["セネガルつええ、ボルト三体くらいいるわ笑笑", \
              "しょーみコロンビアより強い", \
              "それなまちがいないわ"]

>>> encoder = JapaneseTextEncoder(corpus)
>>> encoder.encode("コロンビア強い")
[18, 20, 2]
>>> encoder.vocab
['<pad>', '<unk>', '</s>', '<s>', 'セネガル', 'つえ', 'え', '、', 'ボルト', '三', '体', 'くらい', ' いる', 'わ', '笑', 'しょ', 'ー', 'み', 'コロンビア', 'より', '強い', 'それ', 'な', 'まちがい', 'ない']
>>> encoder.decode(encoder.encode("コロンビア強い"))
コロンビア強い</s>

使い方

1. 日本語の文字列が入ったリストを用意します。

corpus = ["セネガルつええ、ボルト三体くらいいるわ笑笑", 
          "しょーみコロンビアより強い", 
          "それなまちがいないわ"]

2. 用意したデータをJapaneseTextEncoderにぶち込みます:)

encoder = JapaneseTextEncoder(corpus)

これだけです。

min_occurencesはデフォルトで１になっています。これを１０にすると、コーパスの中で１０回以上出現しない単語は未知語となり辞書に登録されません。

encoder = JapaneseTextEncoder(corpus, min_occurences=10)

今はこれぐらいの機能しかありませんが、これから時間をかけて便利なものにしていきたいと考えています。

vocab

辞書に登録されている単語をリストで返す。

word2id

辞書に登録されている単語を｛単語：単語ID｝の辞書形式で返す。

id2word

辞書に登録されている単語を｛単語ID:単語｝の辞書形式で返す。

encode

テキストを単語IDに変換し、リストで返す。

decode

単語IDを受け取って、単語に変換する。

ソースコード

reserved_tokens.py

PADDING_INDEX = 0
UNKNOWN_INDEX = 1
EOS_INDEX = 2
SOS_INDEX = 3
PADDING_TOKEN = '<pad>'
UNKNOWN_TOKEN = '<unk>'
EOS_TOKEN = '</s>'
SOS_TOKEN = '<s>'
RESERVED_ITOS = [PADDING_TOKEN, UNKNOWN_TOKEN, EOS_TOKEN, SOS_TOKEN]
RESERVED_STOI = {token: index for index, token in enumerate(RESERVED_ITOS)}

text_encoder.py

from collections import Counter
from reserved_tokens import SOS_INDEX
from reserved_tokens import EOS_INDEX
from reserved_tokens import UNKNOWN_INDEX
from reserved_tokens import RESERVED_ITOS


class JapaneseTextEncoder:
    """ Encodes the text using a tokenizer.

    Args:
        corpus (list of strings): Text strings to build dictionary on.
        min_occurrences (int, optional): Minimum number of occurences for a token to be
            added to dictionary.
        append_sos (bool, optional): If 'True' append SOS token onto the begin to the encoded vector.
        append_eos (bool, optional): If 'True' append EOS token onto the end to the encoded vector.
        reserved_tokens (list of str, optional): Tokens added to dictionary; reserving the first
            'len(reserved_tokens') indices.

    Example:
        >>> corpus = ["セネガルつええ、ボルト三体くらいいるわ笑笑", \
                      "しょーみコロンビアより強い", \
                      "それなまちがいないわ"]

        >>> encoder = JapaneseTextEncoder(corpus)
        >>> encoder.encode("コロンビア強い")
        [18, 20, 2]
        >>> encoder.vocab
        ['<pad>', '<unk>', '</s>', '<s>', 'セネガル', 'つえ', 'え', '、', 'ボルト', '三', '体', 'くらい', ' いる', 'わ', '笑', 'しょ', 'ー', 'み', 'コロンビア', 'より', '強い', 'それ', 'な', 'まちがい', 'ない']
        >>> encoder.decode(encoder.encode("コロンビア強い"))
        コロンビア強い</s>

    """
    def __init__(self,
                 corpus,
                 min_occurrences=1,
                 append_sos=False,
                 append_eos=True,
                 reserved_tokens=RESERVED_ITOS):
        try:
            import MeCab
        except ImportError:
            print("Please install MeCab.")
            raise

        if not isinstance(corpus, list):
            raise TypeError("Corpus must be a list of strings.")

        self.tagger = MeCab.Tagger("-Owakati")
        self.append_sos = append_sos
        self.append_eos = append_eos
        self.tokens = Counter()

        for sentence in corpus:
            self.tokens.update(self.tokenize(sentence))

        self.itos = reserved_tokens.copy()
        self.stoi = {token: index for index, token in enumerate(reserved_tokens)}
        for token, cnt in self.tokens.items():
            if cnt >= min_occurrences:
                self.itos.append(token)
                self.stoi[token] = len(self.itos) - 1

    @property
    def vocab(self):
        return self.itos

    @property
    def word2id(self):
        return self.stoi

    @property
    def id2word(self):
        return {index: token for token, index in self.stoi.items()}

    def encode(self, sentence, sos_index=SOS_INDEX, eos_index=EOS_INDEX, unknown_index=UNKNOWN_INDEX):
        tokens = self.tokenize(sentence)
        indices = [self.stoi.get(token, unknown_index) for token in tokens]
        if self.append_sos:
            indices.insert(0, sos_index)
        if self.append_eos:
            indices.append(eos_index)
        return indices

    def decode(self, indices):
        tokens = [self.itos[index] for index in indices]
        return "".join(tokens)

    def tokenize(self, sentence):
        return self.tagger.parse(sentence).strip().split()

sample.py

from text_encoder import JapaneseTextEncoder

corpus = ["セネガルつええ、ボルト三体くらいいるわ笑笑",
          "しょーみコロンビアより強い",
          "それなまちがいないわ"]

encoder = JapaneseTextEncoder(corpus)
print(encoder.vocab)
print(encoder.encode("コロンビア強い"))
print(encoder.decode(encoder.encode("コロンビア強い")))
print(encoder.word2id)
print(encoder.id2word)