# AUTOGENERATED! DO NOT EDIT! File to edit: source_nbs/99_test_base.ipynb (unless otherwise specified).

__all__ = ['TestBase', 'PysparkTestBase', 'create_dummy_features_hidden_features', 'test_top_layer']

# Cell

import shutil
import tempfile

import numpy as np
import tensorflow as tf
from loguru import logger
from .input_fn import train_eval_input_fn
from .model_fn import BertMultiTask, BertMultiTaskBody, BertMultiTaskTop
from .params import Params
from .predefined_problems import (
    get_fake_contrastive_learning_fn,
    get_weibo_fake_cls_fn,
    get_weibo_fake_multi_cls_fn,
    get_weibo_fake_ner_fn,
    get_weibo_fake_regression,
    get_weibo_fake_vector_fit,
    get_weibo_masklm, get_weibo_premask_mlm,
    pyspark_fake_cls, pyspark_fake_multi_cls,
    pyspark_fake_seq_tag)
from .special_tokens import EVAL, PREDICT, TRAIN
from .utils import (get_embedding_table_from_model,
                        set_is_pyspark, set_phase)


class TestBase():
    def __init__(self):
        self.setUp()

    def setUp(self) -> None:
        self.tmpfiledir = tempfile.mkdtemp()
        self.tmpckptdir = tempfile.mkdtemp()
        self.prepare_params()

    def tearDown(self) -> None:
        shutil.rmtree(self.tmpfiledir)
        shutil.rmtree(self.tmpckptdir)

    def prepare_params(self):

        self.problem_type_dict = {
            'weibo_fake_ner': 'seq_tag',
            # 'weibo_cws': 'seq_tag',
            'weibo_fake_multi_cls': 'multi_cls',
            'weibo_fake_cls': 'cls',
            'weibo_masklm': 'masklm',
            'weibo_fake_regression': 'regression',
            'weibo_fake_vector_fit': 'vector_fit',
            'weibo_premask_mlm': 'premask_mlm',
            'fake_contrastive_learning': 'contrastive_learning'
        }

        self.processing_fn_dict = {
            'weibo_fake_ner': get_weibo_fake_ner_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            # 'weibo_cws': get_weibo_cws_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_cls': get_weibo_fake_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_multi_cls': get_weibo_fake_multi_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_masklm': get_weibo_masklm(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_regression': get_weibo_fake_regression(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_vector_fit': get_weibo_fake_vector_fit(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_premask_mlm': get_weibo_premask_mlm(file_path='/data/m3tl/data/ner/weiboNER*'),
            'fake_contrastive_learning': get_fake_contrastive_learning_fn(file_path='/data/m3tl/data/ner/weiboNER*')
        }
        self.params = Params()
        self.params.tmp_file_dir = self.tmpfiledir
        self.params.ckpt_dir = self.tmpckptdir
        self.params.transformer_model_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_config_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_model_loading = 'TFAlbertModel'
        self.params.transformer_tokenizer_loading = 'BertTokenizer'
        self.params.transformer_config_loading = 'AlbertConfig'
        self.params.print_data = False
        self.params.enable_modal_type = True
        self.params.num_cpus = 1

        self.params.register_multiple_problems(
            problem_type_dict=self.problem_type_dict, processing_fn_dict=self.processing_fn_dict)
        self.params.assign_problem(
            'weibo_fake_ner&weibo_fake_cls&weibo_fake_regression&weibo_fake_vector_fit|weibo_fake_multi_cls|weibo_masklm|weibo_premask_mlm|fake_contrastive_learning', base_dir=self.tmpckptdir)
        self.params.assign_embedding_layer('duplicate_data_augmentation_embedding')
        self.params.duplicate_data_aug_problems = ['fake_contrastive_learning']
        set_phase(TRAIN)

    def build_body(self):
        if not hasattr(self, 'params'):
            self.prepare_params()
        self.test_body = BertMultiTaskBody(params=self.params)

    def build_top(self):
        if not hasattr(self, 'params'):
            self.prepare_params()
        if not hasattr(self, 'test_body'):
            self.build_body()
        input_embeddings = get_embedding_table_from_model(
            self.test_body.bert.bert_model)
        self.test_top = BertMultiTaskTop(
            params=self.params, input_embeddings=input_embeddings)

    def build_all_model(self):
        if not hasattr(self, 'params'):
            self.prepare_params()
        self.all_model = BertMultiTask(params=self.params)

    def get_dataset(self, mode=TRAIN):
        if not hasattr(self, 'params'):
            self.prepare_params()
        dataset = train_eval_input_fn(
            params=self.params, mode=TRAIN)
        return dataset

    def get_one_batch_input(self, mode=TRAIN):
        dataset = self.get_dataset(mode=mode)
        dataset = dataset.repeat()

        one_batch_data = next(dataset.as_numpy_iterator())
        return one_batch_data

    def get_one_batch_body_model_output(self, mode=TRAIN):
        one_batch_data = self.get_one_batch_input(mode=mode)
        if not hasattr(self, 'test_body'):
            self.build_body()
        features, hidden_features = self.test_body(one_batch_data)
        return features, hidden_features

    def test_mtl_model(self, mtl_model: tf.keras.Model, include_top: bool, mode: str):

        one_batch_body_output = self.get_one_batch_body_model_output(mode=mode)
        if not include_top:
            if not hasattr(self, 'test_top'):
                self.build_top()
            model_outputs = mtl_model(one_batch_body_output)
            return self.test_top(model_outputs)
        else:
            return mtl_model(one_batch_body_output)

    def test_loss_combination_strategy(self, loss_combination_strategy_name: str, run_fit=True):
        if not hasattr(self, 'params'):
            self.prepare_params()
        self.params.assign_loss_combination_strategy(
            loss_combination_strategy_name)

        train_dataset = self.get_dataset(mode=TRAIN)
        eval_dataset = self.get_dataset(mode=EVAL)
        self.build_all_model()

        self.all_model.compile()
        if run_fit:
            self.all_model.fit(
                train_dataset,
                validation_data=eval_dataset,
                steps_per_epoch=1,
                validation_steps=1,
                epochs=2)

    def create_fake_embedding_layer(self) -> tf.keras.layers.Embedding:
        layer = tf.keras.layers.Embedding(input_dim=25000, output_dim=64)
        layer.build((1, 1))
        return layer

    def test_embedding_layer(self, embedding_layer: tf.keras.Model, mode=TRAIN):
        one_batch_data = self.get_one_batch_input(mode=mode)
        return embedding_layer(one_batch_data, mode == TRAIN)


# Cell
class PysparkTestBase(TestBase):
    def __init__(self):
        self.setUp()

    def setUp(self) -> None:
        set_is_pyspark(True)
        self.tmpfiledir = tempfile.mkdtemp()
        self.tmpckptdir = tempfile.mkdtemp()
        self.tmppysparkdir = tempfile.mkdtemp()
        self.prepare_params()

    def prepare_params(self):

        self.problem_type_dict = {
            'pyspark_fake_seq_tag': 'seq_tag',
            'pyspark_fake_multi_cls': 'multi_cls',
            'pyspark_fake_cls': 'cls',
        }

        self.processing_fn_dict = {
            'pyspark_fake_seq_tag': pyspark_fake_seq_tag,
            'pyspark_fake_multi_cls': pyspark_fake_multi_cls,
            'pyspark_fake_cls': pyspark_fake_cls,
        }
        self.params = Params()
        self.params.tmp_file_dir = self.tmpfiledir
        self.params.ckpt_dir = self.tmpckptdir
        self.params.transformer_model_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_config_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_model_loading = 'TFAlbertModel'
        self.params.transformer_tokenizer_loading = 'BertTokenizer'
        self.params.transformer_config_loading = 'AlbertConfig'
        self.params.print_data = False
        self.params.enable_modal_type = True
        self.params.num_cpus = 1
        self.params.pyspark_output_path = self.tmppysparkdir

        self.params.register_multiple_problems(
            problem_type_dict=self.problem_type_dict, processing_fn_dict=self.processing_fn_dict)
        self.params.assign_problem(
            'pyspark_fake_seq_tag&pyspark_fake_multi_cls&pyspark_fake_cls', base_dir=self.tmpckptdir)
        set_phase(TRAIN)


# Cell
# prepare top dummy inputs
# non-empty inputs
def create_dummy_features_hidden_features(batch_size=1, hidden_dim=768, sample_features: dict = None, problem: str = None):
    seq_len = sample_features['text_input_ids'].shape[1]
    dummy_features = {}
    for feature_name, feature in sample_features.items():
        if len(feature.shape) > 1:
            dummy_features[feature_name] = np.zeros(
                shape=(batch_size, *feature.shape[1:]), dtype=feature.dtype)
        else:
            dummy_features[feature_name] = np.zeros(
                shape=(batch_size), dtype=feature.dtype)
    dummy_features['model_input_mask'] = np.ones_like(
        dummy_features['text_input_ids'], dtype=np.int32)

    dummy_features['{}_masked_lm_positions'.format(problem)] = np.zeros(
        shape=(batch_size, 20), dtype=np.int64)
    dummy_features['{}_masked_lm_ids'.format(problem)] = np.zeros(
        shape=(batch_size, 20), dtype=np.int64)
    dummy_features['{}_masked_lm_weights'.format(problem)] = np.ones(
        shape=(batch_size, 20), dtype=np.int64)

    dummy_hidden = {
        'seq': np.zeros(shape=(batch_size, seq_len, hidden_dim), dtype=np.float32),
        'pooled': np.zeros(shape=(batch_size, hidden_dim), dtype=np.float32),
    }
    return dummy_features, dummy_hidden


def test_top_layer(top_class, problem: str, params: Params, sample_features: dict, hidden_dim: int, test_batch_size_list : list=None, **kwargs):
    logger.debug('Testing {}'.format(top_class.__name__))
    top_layer = top_class(params, problem, **kwargs)
    if test_batch_size_list is None:
        test_batch_size_list = [0,1,2]

    for batch_size in test_batch_size_list:
        logger.debug('testing batch size {}'.format(batch_size))
        # non empty test
        d_f, d_h = create_dummy_features_hidden_features(
            batch_size=batch_size, hidden_dim=hidden_dim, sample_features=sample_features, problem=problem)
        set_phase(TRAIN)
        _ = top_layer((d_f, d_h))
        set_phase(EVAL)
        _ = top_layer((d_f, d_h))
        set_phase(PREDICT)
        _ = top_layer((d_f, d_h))

# AUTOGENERATED! DO NOT EDIT! File to edit: source_nbs/99_test_base.ipynb (unless otherwise specified).

__all__ = ['TestBase', 'PysparkTestBase', 'create_dummy_features_hidden_features', 'test_top_layer']

# Cell

import shutil
import tempfile

import numpy as np
import tensorflow as tf
from loguru import logger
from .input_fn import train_eval_input_fn
from .model_fn import BertMultiTask, BertMultiTaskBody, BertMultiTaskTop
from .params import Params
from .predefined_problems import (
    get_fake_contrastive_learning_fn,
    get_weibo_fake_cls_fn,
    get_weibo_fake_multi_cls_fn,
    get_weibo_fake_ner_fn,
    get_weibo_fake_regression,
    get_weibo_fake_vector_fit,
    get_weibo_masklm, get_weibo_premask_mlm,
    pyspark_fake_cls, pyspark_fake_multi_cls,
    pyspark_fake_seq_tag)
from .special_tokens import EVAL, PREDICT, TRAIN
from .utils import (get_embedding_table_from_model,
                        set_is_pyspark, set_phase)


class TestBase():
    def __init__(self):
        self.setUp()

    def setUp(self) -> None:
        self.tmpfiledir = tempfile.mkdtemp()
        self.tmpckptdir = tempfile.mkdtemp()
        self.prepare_params()

    def tearDown(self) -> None:
        shutil.rmtree(self.tmpfiledir)
        shutil.rmtree(self.tmpckptdir)

    def prepare_params(self):

        self.problem_type_dict = {
            'weibo_fake_ner': 'seq_tag',
            # 'weibo_cws': 'seq_tag',
            'weibo_fake_multi_cls': 'multi_cls',
            'weibo_fake_cls': 'cls',
            'weibo_masklm': 'masklm',
            'weibo_fake_regression': 'regression',
            'weibo_fake_vector_fit': 'vector_fit',
            'weibo_premask_mlm': 'premask_mlm',
            'fake_contrastive_learning': 'contrastive_learning'
        }

        self.processing_fn_dict = {
            'weibo_fake_ner': get_weibo_fake_ner_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            # 'weibo_cws': get_weibo_cws_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_cls': get_weibo_fake_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_multi_cls': get_weibo_fake_multi_cls_fn(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_masklm': get_weibo_masklm(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_regression': get_weibo_fake_regression(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_fake_vector_fit': get_weibo_fake_vector_fit(file_path='/data/m3tl/data/ner/weiboNER*'),
            'weibo_premask_mlm': get_weibo_premask_mlm(file_path='/data/m3tl/data/ner/weiboNER*'),
            'fake_contrastive_learning': get_fake_contrastive_learning_fn(file_path='/data/m3tl/data/ner/weiboNER*')
        }
        self.params = Params()
        self.params.tmp_file_dir = self.tmpfiledir
        self.params.ckpt_dir = self.tmpckptdir
        self.params.transformer_model_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_config_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_model_loading = 'TFAlbertModel'
        self.params.transformer_tokenizer_loading = 'BertTokenizer'
        self.params.transformer_config_loading = 'AlbertConfig'
        self.params.print_data = False
        self.params.enable_modal_type = True
        self.params.num_cpus = 1

        self.params.register_multiple_problems(
            problem_type_dict=self.problem_type_dict, processing_fn_dict=self.processing_fn_dict)
        self.params.assign_problem(
            'weibo_fake_ner&weibo_fake_cls&weibo_fake_regression&weibo_fake_vector_fit|weibo_fake_multi_cls|weibo_masklm|weibo_premask_mlm|fake_contrastive_learning', base_dir=self.tmpckptdir)
        self.params.assign_embedding_layer('duplicate_data_augmentation_embedding')
        self.params.duplicate_data_aug_problems = ['fake_contrastive_learning']
        set_phase(TRAIN)

    def build_body(self):
        if not hasattr(self, 'params'):
            self.prepare_params()
        self.test_body = BertMultiTaskBody(params=self.params)

    def build_top(self):
        if not hasattr(self, 'params'):
            self.prepare_params()
        if not hasattr(self, 'test_body'):
            self.build_body()
        input_embeddings = get_embedding_table_from_model(
            self.test_body.bert.bert_model)
        self.test_top = BertMultiTaskTop(
            params=self.params, input_embeddings=input_embeddings)

    def build_all_model(self):
        if not hasattr(self, 'params'):
            self.prepare_params()
        self.all_model = BertMultiTask(params=self.params)

    def get_dataset(self, mode=TRAIN):
        if not hasattr(self, 'params'):
            self.prepare_params()
        dataset = train_eval_input_fn(
            params=self.params, mode=TRAIN)
        return dataset

    def get_one_batch_input(self, mode=TRAIN):
        dataset = self.get_dataset(mode=mode)
        dataset = dataset.repeat()

        one_batch_data = next(dataset.as_numpy_iterator())
        return one_batch_data

    def get_one_batch_body_model_output(self, mode=TRAIN):
        one_batch_data = self.get_one_batch_input(mode=mode)
        if not hasattr(self, 'test_body'):
            self.build_body()
        features, hidden_features = self.test_body(one_batch_data)
        return features, hidden_features

    def test_mtl_model(self, mtl_model: tf.keras.Model, include_top: bool, mode: str):

        one_batch_body_output = self.get_one_batch_body_model_output(mode=mode)
        if not include_top:
            if not hasattr(self, 'test_top'):
                self.build_top()
            model_outputs = mtl_model(one_batch_body_output)
            return self.test_top(model_outputs)
        else:
            return mtl_model(one_batch_body_output)

    def test_loss_combination_strategy(self, loss_combination_strategy_name: str, run_fit=True):
        if not hasattr(self, 'params'):
            self.prepare_params()
        self.params.assign_loss_combination_strategy(
            loss_combination_strategy_name)

        train_dataset = self.get_dataset(mode=TRAIN)
        eval_dataset = self.get_dataset(mode=EVAL)
        self.build_all_model()

        self.all_model.compile()
        if run_fit:
            self.all_model.fit(
                train_dataset,
                validation_data=eval_dataset,
                steps_per_epoch=1,
                validation_steps=1,
                epochs=2)

    def create_fake_embedding_layer(self) -> tf.keras.layers.Embedding:
        layer = tf.keras.layers.Embedding(input_dim=25000, output_dim=64)
        layer.build((1, 1))
        return layer

    def test_embedding_layer(self, embedding_layer: tf.keras.Model, mode=TRAIN):
        one_batch_data = self.get_one_batch_input(mode=mode)
        return embedding_layer(one_batch_data, mode == TRAIN)


# Cell
class PysparkTestBase(TestBase):
    def __init__(self):
        self.setUp()

    def setUp(self) -> None:
        set_is_pyspark(True)
        self.tmpfiledir = tempfile.mkdtemp()
        self.tmpckptdir = tempfile.mkdtemp()
        self.tmppysparkdir = tempfile.mkdtemp()
        self.prepare_params()

    def prepare_params(self):

        self.problem_type_dict = {
            'pyspark_fake_seq_tag': 'seq_tag',
            'pyspark_fake_multi_cls': 'multi_cls',
            'pyspark_fake_cls': 'cls',
        }

        self.processing_fn_dict = {
            'pyspark_fake_seq_tag': pyspark_fake_seq_tag,
            'pyspark_fake_multi_cls': pyspark_fake_multi_cls,
            'pyspark_fake_cls': pyspark_fake_cls,
        }
        self.params = Params()
        self.params.tmp_file_dir = self.tmpfiledir
        self.params.ckpt_dir = self.tmpckptdir
        self.params.transformer_model_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_config_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_tokenizer_name = 'voidful/albert_chinese_tiny'
        self.params.transformer_model_loading = 'TFAlbertModel'
        self.params.transformer_tokenizer_loading = 'BertTokenizer'
        self.params.transformer_config_loading = 'AlbertConfig'
        self.params.print_data = False
        self.params.enable_modal_type = True
        self.params.num_cpus = 1
        self.params.pyspark_output_path = self.tmppysparkdir

        self.params.register_multiple_problems(
            problem_type_dict=self.problem_type_dict, processing_fn_dict=self.processing_fn_dict)
        self.params.assign_problem(
            'pyspark_fake_seq_tag&pyspark_fake_multi_cls&pyspark_fake_cls', base_dir=self.tmpckptdir)
        set_phase(TRAIN)


# Cell
# prepare top dummy inputs
# non-empty inputs
def create_dummy_features_hidden_features(batch_size=1, hidden_dim=768, sample_features: dict = None, problem: str = None):
    seq_len = sample_features['text_input_ids'].shape[1]
    dummy_features = {}
    for feature_name, feature in sample_features.items():
        if len(feature.shape) > 1:
            dummy_features[feature_name] = np.zeros(
                shape=(batch_size, *feature.shape[1:]), dtype=feature.dtype)
        else:
            dummy_features[feature_name] = np.zeros(
                shape=(batch_size), dtype=feature.dtype)
    dummy_features['model_input_mask'] = np.ones_like(
        dummy_features['text_input_ids'], dtype=np.int32)

    dummy_features['{}_masked_lm_positions'.format(problem)] = np.zeros(
        shape=(batch_size, 20), dtype=np.int64)
    dummy_features['{}_masked_lm_ids'.format(problem)] = np.zeros(
        shape=(batch_size, 20), dtype=np.int64)
    dummy_features['{}_masked_lm_weights'.format(problem)] = np.ones(
        shape=(batch_size, 20), dtype=np.int64)

    dummy_hidden = {
        'seq': np.zeros(shape=(batch_size, seq_len, hidden_dim), dtype=np.float32),
        'pooled': np.zeros(shape=(batch_size, hidden_dim), dtype=np.float32),
    }
    return dummy_features, dummy_hidden


def test_top_layer(top_class, problem: str, params: Params, sample_features: dict, hidden_dim: int, test_batch_size_list : list=None, **kwargs):
    logger.debug('Testing {}'.format(top_class.__name__))
    top_layer = top_class(params, problem, **kwargs)
    if test_batch_size_list is None:
        test_batch_size_list = [0,1,2]

    for batch_size in test_batch_size_list:
        logger.debug('testing batch size {}'.format(batch_size))
        # non empty test
        d_f, d_h = create_dummy_features_hidden_features(
            batch_size=batch_size, hidden_dim=hidden_dim, sample_features=sample_features, problem=problem)
        set_phase(TRAIN)
        _ = top_layer((d_f, d_h))
        set_phase(EVAL)
        _ = top_layer((d_f, d_h))
        set_phase(PREDICT)
        _ = top_layer((d_f, d_h))

# AUTOGENERATED! DO NOT EDIT! File to edit: source_nbs/99_test_base.ipynb (unless otherwise specified). # Cell # 'weibo_cws': 'seq_tag', # 'weibo_cws': get_weibo_cws_fn(file_path='/data/m3tl/data/ner/weiboNER*'), # Cell # Cell # prepare top dummy inputs # non-empty inputs # non empty test