<reponame>eladhoffer/seq2seq.pytorch<filename>seq2seq/tools/tokenizer.py<gh_stars>100-1000
# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import os
import string
import codecs
import logging
import tempfile
import sys
from collections import Counter, OrderedDict
import torch
from .config import *

try:
    from subword_nmt import learn_bpe, apply_bpe
    _BPE_AVAILABLE = True
except:
    _BPE_AVAILABLE = False

try:
    from sacremoses import MosesTokenizer, MosesDetokenizer
    _MOSES_AVAILABLE = True
except ImportError:
    _MOSES_AVAILABLE = False

try:
    import sentencepiece as spm
    _SENTENCEPIECE_AVAILABLE = True
except ImportError:
    _SENTENCEPIECE_AVAILABLE = False


class OrderedCounter(Counter, OrderedDict):
    pass


def _segment_words(line, pre_apply=None):
    if pre_apply is not None:
        line = pre_apply(line)
    line = str(line)
    return line.strip('\r\n ').split()


def _get_vocabulary(item_list, segment=_segment_words, from_filenames=True):
    vocab = OrderedCounter()
    if from_filenames:
        filenames = item_list
        # get combined vocabulary of all input files
        for fname in filenames:
            with codecs.open(fname, encoding='UTF-8') as f:
                for line in f:
                    for word in segment(line):
                        vocab[word] += 1
    else:
        for line in item_list:
            for word in segment(line):
                vocab[word] += 1
    return vocab


def _get_double_vocabulary(item_list, segment_words=_segment_words, segment_chars=lambda w: list(w.strip()), from_filenames=True):
    vocab_words = OrderedCounter()
    vocab_chars = OrderedCounter()
    if from_filenames:
        filenames = item_list
        # get combined vocabulary of all input files
        for fname in filenames:
            with codecs.open(fname, encoding='UTF-8') as f:
                for line in f:
                    for word in segment_words(line):
                        vocab_words[word] += 1
                        for char in segment_chars(word):
                            vocab_chars[char] += 1
    else:
        for line in item_list:
            for word in segment_words(line):
                vocab_words[word] += 1
                for char in segment_chars(word):
                    vocab_chars[char] += 1
    return vocab_words, vocab_chars


class Tokenizer(object):

    def __init__(self, file_prefix=None, additional_tokens=None, use_moses=None):
        self.special_tokens = [PAD_TOKEN, UNK_TOKEN, BOS_TOKEN, EOS_TOKEN]
        if file_prefix is not None:
            self.vocab_file = self._vocab_filename(file_prefix)
        if use_moses is not None:
            self.enable_moses(lang=use_moses)
        if additional_tokens is not None:
            self.special_tokens += additional_tokens
        self.__word2idx = {}
        if os.path.isfile(self.vocab_file):
            self.load_vocab(file_prefix)

    def enable_moses(self, lang='en', tokenize=True, detokenize=True):
        if tokenize:
            self._moses_tok = MosesTokenizer(lang=lang)
        else:
            self._moses_tok = None

        if detokenize:
            self._moses_detok = MosesDetokenizer(lang=lang)
        else:
            self._moses_detok = None

    @staticmethod
    def _vocab_filename(model_prefix):
        return '{}.vocab'.format(model_prefix)

    @property
    def vocab_size(self):
        return len(self.vocab) + len(self.special_tokens)

    def pre_tokenize(self, line):
        if hasattr(self, '_moses_tok'):
            return self._moses_tok.tokenize(line, return_str=True)
        return line

    def post_detokenize(self, line):
        if hasattr(self, '_moses_detok'):
            tokens = _segment_words(line)
            return self._moses_detok.detokenize(tokens, return_str=True)
        return line

    def idx2word(self, idx):
        if idx < len(self.special_tokens):
            return self.special_tokens[idx]
        else:
            return self.vocab[idx - len(self.special_tokens)][0]

    def update_word2idx(self):
        self.__word2idx = {
            word[0]: idx + len(self.special_tokens) for idx, word in enumerate(self.vocab)}
        for i, tok in enumerate(self.special_tokens):
            self.__word2idx[tok] = i

    def word2idx(self, word):
        return self.__word2idx.get(word, UNK)

    def segment(self, line, sample=None):
        """segments a line to tokenizable items"""
        line = self.pre_tokenize(line)
        return _segment_words(line)

    def get_vocab(self, item_list, from_filenames=True, limit=None):
        vocab = _get_vocabulary(item_list=item_list, segment=self.segment,
                                from_filenames=from_filenames)
        self.vocab = vocab.most_common(limit)
        self.update_word2idx()

    def save_vocab(self, file_prefix):
        vocab_file = self._vocab_filename(file_prefix)
        if self.vocab is not None:
            with codecs.open(vocab_file, 'w', encoding='UTF-8') as f:
                for (key, freq) in self.vocab:
                    f.write("{0} {1}\n".format(key, freq))

    def load_vocab(self, file_prefix, limit=None, min_count=1):
        vocab_file = self._vocab_filename(file_prefix)
        vocab = OrderedCounter()
        with codecs.open(vocab_file, encoding='UTF-8') as f:
            for line in f:
                try:
                    word, count = line.strip().split()
                except:  # no count
                    word, count = line.strip(), 1
                count = int(count)
                if count >= min_count:
                    vocab[word] = count
        self.vocab = vocab.most_common(limit)
        self.update_word2idx()

    def tokenize(self, line, insert_start=None, insert_end=None, sample=None):
        """tokenize a line, insert_start and insert_end are lists of tokens"""
        inputs = self.segment(line)
        targets = []
        if insert_start is not None:
            targets += insert_start
        for w in inputs:
            targets.append(self.word2idx(w))
        if insert_end is not None:
            targets += insert_end
        return torch.LongTensor(targets)

    def detokenize(self, inputs, delimiter=u' '):
        token_list = [self.idx2word(int(idx)) for idx in inputs]
        outputs = delimiter.join(token_list)
        outputs = self.post_detokenize(outputs)
        return outputs


class BPETokenizer(Tokenizer):

    def __init__(self, file_prefix, additional_tokens=None,
                 num_symbols=10000, min_frequency=2, total_symbols=False, separator='@@',
                 **kwargs):
        super(BPETokenizer, self).__init__(file_prefix=file_prefix,
                                           additional_tokens=additional_tokens,
                                           **kwargs)
        self.num_symbols = num_symbols
        self.min_frequency = min_frequency
        self.total_symbols = total_symbols
        self.separator = separator
        self.vocab_file = self._vocab_filename(file_prefix)
        self.codes_file = '{}.codes'.format(file_prefix)
        if os.path.isfile(self.codes_file):
            self.set_bpe(self.codes_file)

    def set_bpe(self, codes_file):
        with codecs.open(self.codes_file, encoding='UTF-8') as codes:
            self.bpe = apply_bpe.BPE(codes, separator=self.separator)

    def segment(self, line):
        if not hasattr(self, 'bpe'):
            raise NameError('Learn bpe first!')
        line = self.pre_tokenize(line)
        return self.bpe.segment(line.strip('\r\n ')).split(' ')

    def learn_bpe(self, item_list, from_filenames=True):
        logging.info('generating bpe codes file. saving to %s' %
                     self.codes_file)

        # get vocabulary at word level (before bpe)
        def segment_words(line): return _segment_words(line, self.pre_tokenize)
        vocab_words = _get_vocabulary(item_list,
                                      from_filenames=from_filenames,
                                      segment=segment_words)
        vocab_list = ['{0} {1}'.format(key, freq)
                      for (key, freq) in vocab_words.items()]
        # learn BPE on combined vocabulary
        with codecs.open(self.codes_file, 'w', encoding='UTF-8') as output:
            learn_bpe.learn_bpe(vocab_list, output, num_symbols=self.num_symbols,
                                min_frequency=self.min_frequency, verbose=False,
                                is_dict=True, total_symbols=self.total_symbols)
        self.set_bpe(self.codes_file)

    def detokenize(self, inputs, delimiter=' '):
        token_list = [self.idx2word(int(idx)) for idx in inputs]
        detok_string = delimiter.join(token_list)
        try:
            detok_string = detok_string.decode('utf-8')
        except:
            pass
        detok_string = detok_string\
            .replace(self.separator + ' ', '')\
            .replace(self.separator, '')
        detok_string = self.post_detokenize(detok_string)
        return detok_string


class CharTokenizer(Tokenizer):

    def segment(self, line):
        line = self.pre_tokenize(line)
        return list(line.strip())

    def detokenize(self, inputs, delimiter=u''):
        return super(CharTokenizer, self).detokenize(inputs, delimiter)


class SentencePiece(Tokenizer):
    def __init__(self, file_prefix, additional_tokens=None,
                 num_symbols=10000, model_type='unigram',
                 character_coverage=1.0, split_by_whitespace=True,
                 use_moses=None):
        assert _SENTENCEPIECE_AVAILABLE
        self.model_prefix = os.path.abspath(file_prefix)
        self.num_symbols = num_symbols
        self.model_type = model_type
        self.character_coverage = character_coverage
        self.split_by_whitespace = split_by_whitespace
        self.model = None
        if use_moses is not None:
            self.enable_moses(lang=use_moses)
        if additional_tokens is not None:
            self.special_tokens += additional_tokens
        if os.path.isfile('{}.model'.format(file_prefix)):
            self.load_model(file_prefix)

    def serialize_model(self, file_prefix):
        with open(file_prefix, 'rb') as f:
            return f.read()

    def deserialize_model(self, model_serialized):
        fd, path = tempfile.mkstemp()
        try:
            with os.fdopen(fd, 'wb') as tmp:
                tmp.write(model_serialized)
            self.load_model(path, add_suffix=False)
        finally:
            os.remove(path)

    def load_model(self, file_prefix, add_suffix=True):
        if add_suffix:
            file_prefix = '{}.model'.format(file_prefix)
        self.model = spm.SentencePieceProcessor()
        self.model.Load(file_prefix)
        self._model_serialized = self.serialize_model(file_prefix)

    def learn_model(self, file_list, **kwargs):
        file_list = ','.join([os.path.abspath(filename)
                              for filename in file_list])
        self.train_sp_model(input=file_list,
                            model_prefix=self.model_prefix,
                            vocab_size=self.num_symbols,
                            character_coverage=self.character_coverage,
                            model_type=self.model_type,
                            split_by_whitespace=self.split_by_whitespace,
                            **kwargs)
        self.load_model(self.model_prefix)

    def tokenize(self, line, insert_start=None, insert_end=None, sample=None):
        """tokenize a line, insert_start and insert_end are lists of tokens"""
        line = self.pre_tokenize(line)
        if sample is None or sample is False:
            targets = self.model.EncodeAsIds(line)
        else:
            sample = sample if isinstance(sample, dict) else {}
            sample.setdefault('nbest_size', 64)
            sample.setdefault('alpha', 0.1)
            targets = self.model.SampleEncodeAsIds(line, **sample)
        if insert_start is not None:
            targets = insert_start + targets
        if insert_end is not None:
            targets += insert_end
        return torch.LongTensor(targets)

    def detokenize(self, inputs):
        outputs = self.model.DecodeIds([int(idx) for idx in inputs])
        outputs = self.post_detokenize(outputs)
        return outputs

    def idx2word(self, idx):
        return self.model.IdToPiece(idx)

    def word2idx(self, word):
        return self.model.PieceToId(word)

    def segment(self, line, sample=None):
        """segments a line to tokenizable items"""
        line = self.pre_tokenize(line)
        if sample is None or sample is False:
            return self.model.EncodeAsPieces(line)
        else:
            sample = sample if isinstance(sample, dict) else {}
            sample.setdefault('nbest_size', 64)
            sample.setdefault('alpha', 0.1)
            return self.model.SampleEncodeAsPieces(line, **sample)

    @property
    def vocab_size(self):
        return len(self.model)

    @staticmethod
    def train_sp_model(**kwargs):
        """possible arguments:
        --input: one-sentence-per-line raw corpus file. You can pass a comma-separated list of files.
        --model_prefix: output model name prefix. <model_name>.model and <model_name>.vocab are generated.
        --vocab_size: vocabulary size, e.g., 8000, 16000, or 32000
        --character_coverage: amount of characters covered by the model
        --model_type: model type. Choose from unigram (default), bpe, char, or word. The input sentence must be pretokenized when using word type.
        """
        kwargs.update({'unk_piece': UNK_TOKEN, 'bos_piece': BOS_TOKEN,
                       'eos_piece': EOS_TOKEN, 'pad_piece': PAD_TOKEN,
                       'unk_id': UNK, 'bos_id': BOS,
                       'eos_id': EOS, 'pad_id': PAD,
                       'unk_surface': UNK_TOKEN,
                       })
        for arg, val in kwargs.items():
            if isinstance(val, bool):
                kwargs[arg] = 'true' if val else 'false'
        config = ' '.join(['--{}={}'.format(name, value)
                           for name, value in kwargs.items() if value is not None])
        spm.SentencePieceTrainer.Train(config)

    def __getstate__(self):
        state = self.__dict__.copy()
        del state['model']
        return state

    def __setstate__(self, newstate):
        self.deserialize_model(newstate['_model_serialized'])


class WordCharTokenizer(Tokenizer):

    def __init__(self, file_prefix, additional_tokens=None):
        self.word_tokenizer = Tokenizer(
            file_prefix='{}.word'.format(file_prefix), additional_tokens=additional_tokens)
        self.char_tokenizer = CharTokenizer(
            file_prefix='{}.char'.format(file_prefix), additional_tokens=additional_tokens)
        self.vocab = self.word_tokenizer.vocab

    @property
    def vocab_size(self):
        return self.word_tokenizer.vocab_size, self.char_tokenizer.vocab_size

    def common_words(self, num=None, sample=None):
        special = [torch.LongTensor([self.char_tokenizer.word2idx(t)])
                   for t in self.char_tokenizer.special_tokens]
        # self.words_freq_list.most_common(num)]
        words = [w for w, _ in self.word_tokenizer.vocab[:num]]

        return special + [self.char_tokenizer.tokenize(word) for word in words]

    def segment(self, line, sample=None):
        """segments a line to tokenizable items"""
        words = self.word_tokenizer.segment(line, sample=sample)
        chars = [self.char_tokenizer.segment(word) for word in words]
        return words, chars

    def get_vocab(self, item_list, from_filenames=True, limit=None):
        vocab_words, vocab_chars = _get_double_vocabulary(item_list=item_list, segment_words=self.word_tokenizer.segment,
                                                          segment_chars=self.char_tokenizer.segment,
                                                          from_filenames=from_filenames)
        self.words_freq_list = vocab_words
        self.word_tokenizer.vocab = vocab_words.most_common(limit)
        self.word_tokenizer.update_word2idx()
        self.char_tokenizer.vocab = vocab_chars.most_common(limit)
        self.char_tokenizer.update_word2idx()
        self.vocab = self.word_tokenizer.vocab

    def save_vocab(self, vocab_filename):
        self.word_tokenizer.save_vocab('{}.word'.format(vocab_filename))
        self.char_tokenizer.save_vocab('{}.char'.format(vocab_filename))

    def load_vocab(self, vocab_filename, limit=None, min_count=1):
        self.word_tokenizer.load_vocab(
            '{}.word'.format(vocab_filename), limit=limit, min_count=min_count)
        self.char_tokenizer.load_vocab(
            '{}.char'.format(vocab_filename), limit=limit, min_count=min_count)
        self.vocab = self.word_tokenizer.vocab

    def tokenize(self, line, insert_start=None, insert_end=None, sample=None):
        """tokenize a line, insert_start and insert_end are lists of tokens"""
        word_inputs = self.word_tokenizer.segment(line, sample=sample)

        target_words = []
        char_tokens = []

        if insert_start is not None:
            target_words += insert_start
            char_tokens += [torch.LongTensor(insert_start)]
        for w in word_inputs:
            target_words.append(self.word_tokenizer.word2idx(w))

        for i, word in enumerate(word_inputs):
            char_tokens.append(
                self.char_tokenizer.tokenize(word, sample=sample))
        if insert_end is not None:
            target_words += insert_end
            char_tokens += [torch.LongTensor(insert_end)]
        word_tokens = torch.LongTensor(target_words)
        return word_tokens, char_tokens

    def detokenize(self, inputs, delimiter=u' '):
        token_list = []
        for word in inputs:
            token_list.append(self.char_tokenizer.detokenize(word))
        token_list = self.post_detokenize(token_list)
        outputs = delimiter.join(token_list)
        return outputs

# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import os
import string
import codecs
import logging
import tempfile
import sys
from collections import Counter, OrderedDict
import torch
from .config import *

try:
    from subword_nmt import learn_bpe, apply_bpe
    _BPE_AVAILABLE = True
except:
    _BPE_AVAILABLE = False

try:
    from sacremoses import MosesTokenizer, MosesDetokenizer
    _MOSES_AVAILABLE = True
except ImportError:
    _MOSES_AVAILABLE = False

try:
    import sentencepiece as spm
    _SENTENCEPIECE_AVAILABLE = True
except ImportError:
    _SENTENCEPIECE_AVAILABLE = False


class OrderedCounter(Counter, OrderedDict):
    pass


def _segment_words(line, pre_apply=None):
    if pre_apply is not None:
        line = pre_apply(line)
    line = str(line)
    return line.strip('\r\n ').split()


def _get_vocabulary(item_list, segment=_segment_words, from_filenames=True):
    vocab = OrderedCounter()
    if from_filenames:
        filenames = item_list
        # get combined vocabulary of all input files
        for fname in filenames:
            with codecs.open(fname, encoding='UTF-8') as f:
                for line in f:
                    for word in segment(line):
                        vocab[word] += 1
    else:
        for line in item_list:
            for word in segment(line):
                vocab[word] += 1
    return vocab


def _get_double_vocabulary(item_list, segment_words=_segment_words, segment_chars=lambda w: list(w.strip()), from_filenames=True):
    vocab_words = OrderedCounter()
    vocab_chars = OrderedCounter()
    if from_filenames:
        filenames = item_list
        # get combined vocabulary of all input files
        for fname in filenames:
            with codecs.open(fname, encoding='UTF-8') as f:
                for line in f:
                    for word in segment_words(line):
                        vocab_words[word] += 1
                        for char in segment_chars(word):
                            vocab_chars[char] += 1
    else:
        for line in item_list:
            for word in segment_words(line):
                vocab_words[word] += 1
                for char in segment_chars(word):
                    vocab_chars[char] += 1
    return vocab_words, vocab_chars


class Tokenizer(object):

    def __init__(self, file_prefix=None, additional_tokens=None, use_moses=None):
        self.special_tokens = [PAD_TOKEN, UNK_TOKEN, BOS_TOKEN, EOS_TOKEN]
        if file_prefix is not None:
            self.vocab_file = self._vocab_filename(file_prefix)
        if use_moses is not None:
            self.enable_moses(lang=use_moses)
        if additional_tokens is not None:
            self.special_tokens += additional_tokens
        self.__word2idx = {}
        if os.path.isfile(self.vocab_file):
            self.load_vocab(file_prefix)

    def enable_moses(self, lang='en', tokenize=True, detokenize=True):
        if tokenize:
            self._moses_tok = MosesTokenizer(lang=lang)
        else:
            self._moses_tok = None

        if detokenize:
            self._moses_detok = MosesDetokenizer(lang=lang)
        else:
            self._moses_detok = None

    @staticmethod
    def _vocab_filename(model_prefix):
        return '{}.vocab'.format(model_prefix)

    @property
    def vocab_size(self):
        return len(self.vocab) + len(self.special_tokens)

    def pre_tokenize(self, line):
        if hasattr(self, '_moses_tok'):
            return self._moses_tok.tokenize(line, return_str=True)
        return line

    def post_detokenize(self, line):
        if hasattr(self, '_moses_detok'):
            tokens = _segment_words(line)
            return self._moses_detok.detokenize(tokens, return_str=True)
        return line

    def idx2word(self, idx):
        if idx < len(self.special_tokens):
            return self.special_tokens[idx]
        else:
            return self.vocab[idx - len(self.special_tokens)][0]

    def update_word2idx(self):
        self.__word2idx = {
            word[0]: idx + len(self.special_tokens) for idx, word in enumerate(self.vocab)}
        for i, tok in enumerate(self.special_tokens):
            self.__word2idx[tok] = i

    def word2idx(self, word):
        return self.__word2idx.get(word, UNK)

    def segment(self, line, sample=None):
        """segments a line to tokenizable items"""
        line = self.pre_tokenize(line)
        return _segment_words(line)

    def get_vocab(self, item_list, from_filenames=True, limit=None):
        vocab = _get_vocabulary(item_list=item_list, segment=self.segment,
                                from_filenames=from_filenames)
        self.vocab = vocab.most_common(limit)
        self.update_word2idx()

    def save_vocab(self, file_prefix):
        vocab_file = self._vocab_filename(file_prefix)
        if self.vocab is not None:
            with codecs.open(vocab_file, 'w', encoding='UTF-8') as f:
                for (key, freq) in self.vocab:
                    f.write("{0} {1}\n".format(key, freq))

    def load_vocab(self, file_prefix, limit=None, min_count=1):
        vocab_file = self._vocab_filename(file_prefix)
        vocab = OrderedCounter()
        with codecs.open(vocab_file, encoding='UTF-8') as f:
            for line in f:
                try:
                    word, count = line.strip().split()
                except:  # no count
                    word, count = line.strip(), 1
                count = int(count)
                if count >= min_count:
                    vocab[word] = count
        self.vocab = vocab.most_common(limit)
        self.update_word2idx()

    def tokenize(self, line, insert_start=None, insert_end=None, sample=None):
        """tokenize a line, insert_start and insert_end are lists of tokens"""
        inputs = self.segment(line)
        targets = []
        if insert_start is not None:
            targets += insert_start
        for w in inputs:
            targets.append(self.word2idx(w))
        if insert_end is not None:
            targets += insert_end
        return torch.LongTensor(targets)

    def detokenize(self, inputs, delimiter=u' '):
        token_list = [self.idx2word(int(idx)) for idx in inputs]
        outputs = delimiter.join(token_list)
        outputs = self.post_detokenize(outputs)
        return outputs


class BPETokenizer(Tokenizer):

    def __init__(self, file_prefix, additional_tokens=None,
                 num_symbols=10000, min_frequency=2, total_symbols=False, separator='@@',
                 **kwargs):
        super(BPETokenizer, self).__init__(file_prefix=file_prefix,
                                           additional_tokens=additional_tokens,
                                           **kwargs)
        self.num_symbols = num_symbols
        self.min_frequency = min_frequency
        self.total_symbols = total_symbols
        self.separator = separator
        self.vocab_file = self._vocab_filename(file_prefix)
        self.codes_file = '{}.codes'.format(file_prefix)
        if os.path.isfile(self.codes_file):
            self.set_bpe(self.codes_file)

    def set_bpe(self, codes_file):
        with codecs.open(self.codes_file, encoding='UTF-8') as codes:
            self.bpe = apply_bpe.BPE(codes, separator=self.separator)

    def segment(self, line):
        if not hasattr(self, 'bpe'):
            raise NameError('Learn bpe first!')
        line = self.pre_tokenize(line)
        return self.bpe.segment(line.strip('\r\n ')).split(' ')

    def learn_bpe(self, item_list, from_filenames=True):
        logging.info('generating bpe codes file. saving to %s' %
                     self.codes_file)

        # get vocabulary at word level (before bpe)
        def segment_words(line): return _segment_words(line, self.pre_tokenize)
        vocab_words = _get_vocabulary(item_list,
                                      from_filenames=from_filenames,
                                      segment=segment_words)
        vocab_list = ['{0} {1}'.format(key, freq)
                      for (key, freq) in vocab_words.items()]
        # learn BPE on combined vocabulary
        with codecs.open(self.codes_file, 'w', encoding='UTF-8') as output:
            learn_bpe.learn_bpe(vocab_list, output, num_symbols=self.num_symbols,
                                min_frequency=self.min_frequency, verbose=False,
                                is_dict=True, total_symbols=self.total_symbols)
        self.set_bpe(self.codes_file)

    def detokenize(self, inputs, delimiter=' '):
        token_list = [self.idx2word(int(idx)) for idx in inputs]
        detok_string = delimiter.join(token_list)
        try:
            detok_string = detok_string.decode('utf-8')
        except:
            pass
        detok_string = detok_string\
            .replace(self.separator + ' ', '')\
            .replace(self.separator, '')
        detok_string = self.post_detokenize(detok_string)
        return detok_string


class CharTokenizer(Tokenizer):

    def segment(self, line):
        line = self.pre_tokenize(line)
        return list(line.strip())

    def detokenize(self, inputs, delimiter=u''):
        return super(CharTokenizer, self).detokenize(inputs, delimiter)


class SentencePiece(Tokenizer):
    def __init__(self, file_prefix, additional_tokens=None,
                 num_symbols=10000, model_type='unigram',
                 character_coverage=1.0, split_by_whitespace=True,
                 use_moses=None):
        assert _SENTENCEPIECE_AVAILABLE
        self.model_prefix = os.path.abspath(file_prefix)
        self.num_symbols = num_symbols
        self.model_type = model_type
        self.character_coverage = character_coverage
        self.split_by_whitespace = split_by_whitespace
        self.model = None
        if use_moses is not None:
            self.enable_moses(lang=use_moses)
        if additional_tokens is not None:
            self.special_tokens += additional_tokens
        if os.path.isfile('{}.model'.format(file_prefix)):
            self.load_model(file_prefix)

    def serialize_model(self, file_prefix):
        with open(file_prefix, 'rb') as f:
            return f.read()

    def deserialize_model(self, model_serialized):
        fd, path = tempfile.mkstemp()
        try:
            with os.fdopen(fd, 'wb') as tmp:
                tmp.write(model_serialized)
            self.load_model(path, add_suffix=False)
        finally:
            os.remove(path)

    def load_model(self, file_prefix, add_suffix=True):
        if add_suffix:
            file_prefix = '{}.model'.format(file_prefix)
        self.model = spm.SentencePieceProcessor()
        self.model.Load(file_prefix)
        self._model_serialized = self.serialize_model(file_prefix)

    def learn_model(self, file_list, **kwargs):
        file_list = ','.join([os.path.abspath(filename)
                              for filename in file_list])
        self.train_sp_model(input=file_list,
                            model_prefix=self.model_prefix,
                            vocab_size=self.num_symbols,
                            character_coverage=self.character_coverage,
                            model_type=self.model_type,
                            split_by_whitespace=self.split_by_whitespace,
                            **kwargs)
        self.load_model(self.model_prefix)

    def tokenize(self, line, insert_start=None, insert_end=None, sample=None):
        """tokenize a line, insert_start and insert_end are lists of tokens"""
        line = self.pre_tokenize(line)
        if sample is None or sample is False:
            targets = self.model.EncodeAsIds(line)
        else:
            sample = sample if isinstance(sample, dict) else {}
            sample.setdefault('nbest_size', 64)
            sample.setdefault('alpha', 0.1)
            targets = self.model.SampleEncodeAsIds(line, **sample)
        if insert_start is not None:
            targets = insert_start + targets
        if insert_end is not None:
            targets += insert_end
        return torch.LongTensor(targets)

    def detokenize(self, inputs):
        outputs = self.model.DecodeIds([int(idx) for idx in inputs])
        outputs = self.post_detokenize(outputs)
        return outputs

    def idx2word(self, idx):
        return self.model.IdToPiece(idx)

    def word2idx(self, word):
        return self.model.PieceToId(word)

    def segment(self, line, sample=None):
        """segments a line to tokenizable items"""
        line = self.pre_tokenize(line)
        if sample is None or sample is False:
            return self.model.EncodeAsPieces(line)
        else:
            sample = sample if isinstance(sample, dict) else {}
            sample.setdefault('nbest_size', 64)
            sample.setdefault('alpha', 0.1)
            return self.model.SampleEncodeAsPieces(line, **sample)

    @property
    def vocab_size(self):
        return len(self.model)

    @staticmethod
    def train_sp_model(**kwargs):
        """possible arguments:
        --input: one-sentence-per-line raw corpus file. You can pass a comma-separated list of files.
        --model_prefix: output model name prefix. <model_name>.model and <model_name>.vocab are generated.
        --vocab_size: vocabulary size, e.g., 8000, 16000, or 32000
        --character_coverage: amount of characters covered by the model
        --model_type: model type. Choose from unigram (default), bpe, char, or word. The input sentence must be pretokenized when using word type.
        """
        kwargs.update({'unk_piece': UNK_TOKEN, 'bos_piece': BOS_TOKEN,
                       'eos_piece': EOS_TOKEN, 'pad_piece': PAD_TOKEN,
                       'unk_id': UNK, 'bos_id': BOS,
                       'eos_id': EOS, 'pad_id': PAD,
                       'unk_surface': UNK_TOKEN,
                       })
        for arg, val in kwargs.items():
            if isinstance(val, bool):
                kwargs[arg] = 'true' if val else 'false'
        config = ' '.join(['--{}={}'.format(name, value)
                           for name, value in kwargs.items() if value is not None])
        spm.SentencePieceTrainer.Train(config)

    def __getstate__(self):
        state = self.__dict__.copy()
        del state['model']
        return state

    def __setstate__(self, newstate):
        self.deserialize_model(newstate['_model_serialized'])


class WordCharTokenizer(Tokenizer):

    def __init__(self, file_prefix, additional_tokens=None):
        self.word_tokenizer = Tokenizer(
            file_prefix='{}.word'.format(file_prefix), additional_tokens=additional_tokens)
        self.char_tokenizer = CharTokenizer(
            file_prefix='{}.char'.format(file_prefix), additional_tokens=additional_tokens)
        self.vocab = self.word_tokenizer.vocab

    @property
    def vocab_size(self):
        return self.word_tokenizer.vocab_size, self.char_tokenizer.vocab_size

    def common_words(self, num=None, sample=None):
        special = [torch.LongTensor([self.char_tokenizer.word2idx(t)])
                   for t in self.char_tokenizer.special_tokens]
        # self.words_freq_list.most_common(num)]
        words = [w for w, _ in self.word_tokenizer.vocab[:num]]

        return special + [self.char_tokenizer.tokenize(word) for word in words]

    def segment(self, line, sample=None):
        """segments a line to tokenizable items"""
        words = self.word_tokenizer.segment(line, sample=sample)
        chars = [self.char_tokenizer.segment(word) for word in words]
        return words, chars

    def get_vocab(self, item_list, from_filenames=True, limit=None):
        vocab_words, vocab_chars = _get_double_vocabulary(item_list=item_list, segment_words=self.word_tokenizer.segment,
                                                          segment_chars=self.char_tokenizer.segment,
                                                          from_filenames=from_filenames)
        self.words_freq_list = vocab_words
        self.word_tokenizer.vocab = vocab_words.most_common(limit)
        self.word_tokenizer.update_word2idx()
        self.char_tokenizer.vocab = vocab_chars.most_common(limit)
        self.char_tokenizer.update_word2idx()
        self.vocab = self.word_tokenizer.vocab

    def save_vocab(self, vocab_filename):
        self.word_tokenizer.save_vocab('{}.word'.format(vocab_filename))
        self.char_tokenizer.save_vocab('{}.char'.format(vocab_filename))

    def load_vocab(self, vocab_filename, limit=None, min_count=1):
        self.word_tokenizer.load_vocab(
            '{}.word'.format(vocab_filename), limit=limit, min_count=min_count)
        self.char_tokenizer.load_vocab(
            '{}.char'.format(vocab_filename), limit=limit, min_count=min_count)
        self.vocab = self.word_tokenizer.vocab

    def tokenize(self, line, insert_start=None, insert_end=None, sample=None):
        """tokenize a line, insert_start and insert_end are lists of tokens"""
        word_inputs = self.word_tokenizer.segment(line, sample=sample)

        target_words = []
        char_tokens = []

        if insert_start is not None:
            target_words += insert_start
            char_tokens += [torch.LongTensor(insert_start)]
        for w in word_inputs:
            target_words.append(self.word_tokenizer.word2idx(w))

        for i, word in enumerate(word_inputs):
            char_tokens.append(
                self.char_tokenizer.tokenize(word, sample=sample))
        if insert_end is not None:
            target_words += insert_end
            char_tokens += [torch.LongTensor(insert_end)]
        word_tokens = torch.LongTensor(target_words)
        return word_tokens, char_tokens

    def detokenize(self, inputs, delimiter=u' '):
        token_list = []
        for word in inputs:
            token_list.append(self.char_tokenizer.detokenize(word))
        token_list = self.post_detokenize(token_list)
        outputs = delimiter.join(token_list)
        return outputs
# -*- coding: utf-8 -*- # get combined vocabulary of all input files # get combined vocabulary of all input files segments a line to tokenizable items # no count tokenize a line, insert_start and insert_end are lists of tokens # get vocabulary at word level (before bpe) # learn BPE on combined vocabulary tokenize a line, insert_start and insert_end are lists of tokens segments a line to tokenizable items possible arguments:         --input: one-sentence-per-line raw corpus file. You can pass a comma-separated list of files.         --model_prefix: output model name prefix. <model_name>.model and <model_name>.vocab are generated.         --vocab_size: vocabulary size, e.g., 8000, 16000, or 32000         --character_coverage: amount of characters covered by the model         --model_type: model type. Choose from unigram (default), bpe, char, or word. The input sentence must be pretokenized when using word type. # self.words_freq_list.most_common(num)] segments a line to tokenizable items tokenize a line, insert_start and insert_end are lists of tokens