multifit/fastai_contrib/text_data.py

import os
import pathlib
import pickle
from functools import reduce
from typing import Collection, List

from pandas import DataFrame
from sacremoses import MosesTokenizer

import fastai
from fastai.basic_data import DataBunch

from fastai.core import ListRules, PathOrStr, defaults, IntsOrStrs, is_listy
from fastai.data_block import ItemLists
from fastai.text import Tokenizer, BaseTokenizer, Vocab, SPProcessor, TextList, TextLMDataBunch


class MosesPreprocessingFunc():
    def __init__(self, lang: str):
        self.mt = MosesTokenizer(lang)

    def __call__(self, t: str) -> str:
        return self.mt.tokenize(t, return_str=True, escape=True)


class SentencePieceTokenizer(Tokenizer):
    "Put together rules and a tokenizer function to tokenize text with multiprocessing."
    def __init__(self, spm_model, lang:str='en', pre_rules:ListRules=None,
                 post_rules:ListRules=None, special_cases:Collection[str]=None, n_cpus:int=None):
        # moses is added to preprocessing functions
        super().__init__(self.tok_fun_with_sp, lang, pre_rules, post_rules, special_cases, n_cpus)
        self.spm_model = spm_model

    def tok_fun_with_sp(self, lang):
        try:
            import sentencepiece as spm
        except ImportError:
            raise Exception('sentencepiece module is missing: run `pip install sentencepiece`')
        tok = BaseTokenizer(lang)
        tok.sp = spm.SentencePieceProcessor()
        tok.sp.Load(str(self.spm_model))
        return tok

    def process_text(self, t:str, tok:BaseTokenizer) -> List[str]:
        "Process one text `t` with tokenizer `tok`."
        toks = super().process_text(t, tok)
        toks = tok.sp.EncodeAsPieces(" ".join(toks))
        return toks


full_char_coverage_langs = ["bg", "cs", "da", "de", "el", "en", "es", "et", "fi", "fr", "ga", "hr", "hu",
                       "it","lt","lv","mt","nl","pl","pt","ro","sk","sl","sv"] # all European langus


def get_sentencepiece(cache_dir:PathOrStr, load_text, pre_rules: ListRules=None, post_rules:ListRules=None,
                      vocab_size:int=30000, model_type:str='unigram', input_sentence_size:int=1E7, lang='en', fixed_character_coverage=False):
    try:
        import sentencepiece as spm
    except ImportError:
        raise Exception('sentencepiece module is missing: run `pip install sentencepiece`')

    cache_dir = pathlib.Path(cache_dir)
    pre_rules = pre_rules if pre_rules is not None else defaults.text_pre_rules
    post_rules = post_rules if post_rules is not None else defaults.text_post_rules

    special_cases = defaults.text_spec_tok # + ['xxlink', 'xxuser', 'xxnumber', 'xxemoji', 'yyemoji']
    if not os.path.isfile(cache_dir / 'spm.model') or not os.path.isfile(cache_dir / f'itos.pkl'):
        # load the text from the train tokens file
        text = load_text()
        text = filter(lambda x: len(x.rstrip(" ")), text)
        text = (reduce(lambda t, rule: rule(t), pre_rules, line) for line in text)
        def cleanup_n_postprocess(t):
            t = t.split()
            for r in post_rules:
                t = r(t)
            return ' '.join(t)
        text = map(cleanup_n_postprocess, text)
        raw_text_path = cache_dir / 'all_text.txt'
        with open(raw_text_path, 'w') as f: f.write("\n".join(text))

        if fixed_character_coverage:
            char_coverage = 0.9995
        else:
            char_coverage = 1 if lang in full_char_coverage_langs else 0.99

        sp_params = [
            f"--input={raw_text_path}",
            f"--character_coverage={char_coverage}",
            f"--unk_id={len(special_cases)}",
            f"--pad_id=-1",
            f"--bos_id=-1",
            f"--eos_id=-1",
            f"--max_sentence_length=20480",
            f"--input_sentence_size={int(input_sentence_size)}",
            f"--user_defined_symbols={','.join(special_cases)}",
            f"--model_prefix={cache_dir/'spm'}",
            f"--vocab_size={vocab_size} --model_type={model_type}"]
        spm.SentencePieceTrainer.Train(" ".join(sp_params))

        with open(cache_dir / 'spm.vocab', 'r') as f:
            vocab = [line.split('\t')[0] for line in f.readlines()]

        pickle.dump(vocab, open(cache_dir/ f'itos.pkl', 'wb'))
    # todo add post rules
    vocab = Vocab(pickle.load(open(cache_dir / f'itos.pkl', 'rb')))
    # We cannot use lambdas or local methods here, since `tok_func` needs to be
    # pickle-able in order to be called in subprocesses when multithread tokenizing
    tokenizer = SentencePieceTokenizer(cache_dir/'spm.model',
                                lang=lang,
                                pre_rules=pre_rules,
                                post_rules=post_rules)
    return {'tokenizer': tokenizer, 'vocab': vocab}

class SPProcessor2(SPProcessor):
    def process(self, ds):
        super().process(ds)
        ds.vocab.sp_model = self.sp_model
        ds.vocab.sp_vocab = self.sp_vocab

def get_sentencepiece_fastai(cache_dir: PathOrStr,  pre_rules: ListRules = None,
                          post_rules: ListRules = None,
                          vocab_size: int = 30000, lang='en'):
    cache_dir = pathlib.Path(cache_dir)

    sp_model = cache_dir / 'spm.model'
    if not sp_model.is_file():
        sp_model = None
    sp_vocab = cache_dir / 'spm.vocab'
    if not sp_vocab.is_file():
        sp_vocab = None
    processor = SPProcessor2(
        pre_rules=pre_rules,
        post_rules=post_rules,
        mark_fields=True,
        vocab_sz=vocab_size,
        sp_model=sp_model,
        sp_vocab=sp_vocab,
        lang=lang,
        tmp_dir=cache_dir.absolute()  # absolute make sure that dataset path is not added as prefix
    )
    return {'processor': processor}

# temporary loading function as from_df does not support processors
def make_data_bunch_from_df(cls, path: PathOrStr, train_df: DataFrame, valid_df: DataFrame,
            tokenizer: Tokenizer = None, vocab: Vocab = None, classes: Collection[str] = None,
            text_cols: IntsOrStrs = 1,
            label_cols: IntsOrStrs = 0, label_delim: str = None, chunksize: int = 10000,
            max_vocab: int = 60000,
            min_freq: int = 2, mark_fields: bool = False, include_bos: bool = True,
            include_eos: bool = False, processor=None, **kwargs) -> DataBunch:
    "Create a `TextDataBunch` from DataFrames. `kwargs` are passed to the dataloader creation."
    assert processor is None or tokenizer is None, "Processor and tokenizer are mutually exclusive."

    if processor is None:
        processor = fastai.text.data._get_processor(tokenizer=tokenizer, vocab=vocab, chunksize=chunksize, max_vocab=max_vocab,
                                   min_freq=min_freq, mark_fields=mark_fields,
                                   include_bos=include_bos, include_eos=include_eos)

    if classes is None and is_listy(label_cols) and len(label_cols) > 1: classes = label_cols
    src = ItemLists(path, TextList.from_df(train_df, path, cols=text_cols, processor=processor),
                    TextList.from_df(valid_df, path, cols=text_cols, processor=processor))
    if cls == TextLMDataBunch:
        src = src.label_for_lm()
    else:
        if label_delim is not None:
            src = src.label_from_df(cols=label_cols, classes=classes, label_delim=label_delim)
        else:
            src = src.label_from_df(cols=label_cols, classes=classes)
    return src.databunch(**kwargs)