PyTorch-NLPを使ってみる - どん底から這い上がるまでの記録

"""
    japanese_encoder.py
"""
from torchnlp.text_encoders.static_tokenizer_encoder import StaticTokenizerEncoder


class JapaneseTokenizer:
    def __init__(self, tagger):
        self.tagger = tagger

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


class JapaneseEncoder(StaticTokenizerEncoder):
    """ Encodes the text written in Japanese by splitting on whitespace.

    Args:
        sample (list of strings): Sample of data to build dictionary on
        min_occurrences (int, optional): Minimum number of occurrences for a token to
        be added to dictionary.

        append_eos (bool, optional): If 'True' append EOS token onto the end to the
        encoded vector.

    Example:

        >>> encoder = JapaneseEncoder(["今日はいい天気ですね。", "そうですね。"])

        >>> encoder.encode("今日はいい天気ですね。")
        tensor([  5,   6,   7,   8,   9,  10,  11])

        >>> encoder.vocab
        ['<pad>', '<unk>', '</s>', '<s>', '<copy>', '今日', 'は', 'いい', '天気', 'です', 'ね', '。', 'そうですね']

        >>> encoder.decode(encoder.encode("今日はいい天気ですね。"))
        '今日はいい天気ですね。'

    """

    def __init__(self, *args, **kwargs):
        if 'tokenize' in kwargs:
            raise TypeError('JapaneseEncoder defines a tokenize callable MeCab')
        try:
            import MeCab
        except ImportError:
            print("Please install MeCab.")
            raise

        self.jatokenizer = JapaneseTokenizer(MeCab.Tagger("-Owakati")).tokenize
        super().__init__(*args, tokenize=self.jatokenizer, **kwargs)

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

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

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

サンプルプログラム①

新しく作ったjapanese_encoder.pyをインポートしています。

from japanese_encoder import JapaneseEncoder


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

encoder = JapaneseEncoder(sample)
sentence = "ポーランド強い"

print(encoder.encode(sentence))
# tensor([  1,  21])

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

print(encoder.decode(encoder.encode(sentence)))
# <unk>強い

<unk>は未知語を表します。

サンプルプログラム②

JapaenseEncoderに渡すオプションを少し変えています。

min_occurrencesは指定した値以上の出現する単語を辞書に加えます。

append_eosをTrueにするとテキストをTensorに変換したときに終端記号を加えるようになります。

encoder = JapaneseEncoder(sample, min_occurrences=2, append_eos=True)
sentence = "ポーランド強い"

print(encoder.encode(sentence))
# tensor([  1,  1, 2])

print(encoder.vocab)
# ['<pad>', '<unk>', '</s>', '<s>', '<copy>', 'わ', '笑']

print(encoder.decode(encoder.encode(sentence)))
# <unk><unk></s>