{"cells":[{"metadata":{},"cell_type":"markdown","source":"Another way of dataloader instead of using tf.Example and tf.Record. (I don't know the difference)\nThis is just a sample and you can't run because of private utility scripts. sorry for inconvenient."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/transformers')\nimport json\nfrom typing import Tuple, List, Generator, Union, Mapping\nfrom pathlib import Path\nimport numpy as np\nimport tensorflow as tf\n\nfrom common_dataloader_base import BaseDataloader\nfrom common_utils_config import Config, Mode\nfrom common_preprocess_base import BasePreprocessor\nfrom common_preprocess_nq_preprocess import NQPreprocessor\nfrom common_entity_nq_entity import Sample, NQSample, Record","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TrainInputType = Tuple[\n    np.int64,    # unique_id\n    np.int64,    # example_id\n    np.ndarray,  # input_ids\n    np.ndarray,  # input_mask\n    np.ndarray,  # segment_ids\n    np.ndarray,  # token_map\n    np.ndarray,  # max_context_map\n    np.int32,    # answer_type\n    np.int32,    # short_start_token\n    np.int32,    # short_end_token\n    np.int32,    # long_start_token\n    np.int32     # long_end_token\n]\n\nTestInputType = Tuple[\n    np.int64,    # unique_id\n    np.int64,    # example_id\n    np.ndarray,  # input_ids\n    np.ndarray,  # input_mask\n    np.ndarray,  # segment_ids\n    np.ndarray,  # token_map\n    np.ndarray   # max_context_map\n]","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class NQDataloader(BaseDataloader):\n\n    def __init__(self,\n                 config: Config,\n                 filepath: Path,\n                 preprocessor: BasePreprocessor,\n                 mode: Mode = Mode.TRAIN,\n                 is_debug: bool = False) -> None:\n        self.config = config\n        self.is_debug = is_debug\n        self.filepath = filepath\n        self.mode = mode\n        self.preprocessor = preprocessor\n\n    def generator(self) -> Generator:\n        '''Return Generator for tf.data.Dataset\n        '''\n        for line in open(self.filepath, 'r'):\n            line = json.loads(line)\n            record = Record(**line)\n\n            samples, nq_sample = self.preprocess(record)\n            for sample in samples:\n                if self.is_debug:\n                    yield sample, nq_sample, record\n                elif self.mode != Mode.TRAIN:\n                    yield self.to_input(sample, self.mode)\n                else:\n                    if sample.long_span is not None or sample.short_span is not None:\n                        yield self.to_input(sample, self.mode)\n                    else:\n                        if np.random.rand() < self.config.downsample_rate:  # = 0.02\n                            yield self.to_input(sample, self.mode)\n\n    def preprocess(self, record: Record) -> Tuple[List[Sample], NQSample]:\n        '''Apply preprocess by Preprocessor\n        '''\n        samples, nq_sample = self.preprocessor.preprocess(record)\n        return samples, nq_sample\n\n    def to_input(self,\n                 sample: Sample,\n                 mode: Mode) -> Union[TrainInputType, TestInputType]:\n        '''Make model input\n        '''\n        unique_id = np.int64(sample.unique_id)\n        example_id = np.int64(sample.example_id)\n        input_ids = np.array(sample.input_ids, dtype=np.int32)\n        input_mask = np.array(sample.input_mask, dtype=np.int32)\n        segment_ids = np.array(sample.segment_ids, dtype=np.int32)\n        token_map = [-1] * len(sample.input_ids)\n        max_context_map = [-1] * len(sample.input_ids)\n        for k, v in sample.wp_to_token_map.items():\n            token_map[k] = v\n        for k, v in sample.wp_token_to_max_context_map.items():\n            max_context_map[k] = v\n        if mode == Mode.TEST:\n            return unique_id, example_id, input_ids, input_mask, segment_ids, np.array(token_map), np.array(max_context_map)\n        else:\n            answer_type = np.int32(sample.answer_type.value)\n            short_start_token = np.int32(sample.short_span.start_token if sample.short_span else -1)\n            short_end_token = np.int32(sample.short_span.end_token if sample.short_span else -1)\n            long_start_token = np.int32(sample.long_span.start_token if sample.long_span else -1)\n            long_end_token = np.int32(sample.long_span.end_token if sample.long_span else -1)\n            return (unique_id, example_id, input_ids, input_mask, segment_ids, np.array(token_map), np.array(max_context_map),\n                    answer_type, short_start_token, short_end_token, long_start_token, long_end_token)\n\n    def to_dataset(self, mode: Mode) -> tf.data.Dataset:\n        '''Return tf.data.Dataset instance for training\n        '''\n        dataset = tf.data.Dataset.from_generator(\n            self.generator,\n            output_types=self.train_output_type if mode != Mode.TEST else self.test_output_type,\n            output_shapes=self.train_output_shape if mode != Mode.TEST else self.test_output_shape\n        )\n        batch_size = self.config.batch_size if mode == Mode.TRAIN else self.config.test_batch_size\n        dataset = dataset.batch(batch_size).prefetch(4)\n        return dataset\n\n    @property\n    def train_output_type(self) -> Tuple:\n        return (\n            tf.int64,  # unique_id\n            tf.int64,  # example_id\n            tf.int32,  # input_ids\n            tf.int32,  # input_mask\n            tf.int32,  # segment_ids\n            tf.int32,  # token_map\n            tf.int32,  # max_context_map\n            tf.int32,  # answer_type\n            tf.int32,  # short_start_token\n            tf.int32,  # short_end_token\n            tf.int32,  # long_start_token\n            tf.int32   # long_end_token\n        )\n\n    @property\n    def train_output_shape(self) -> Tuple:\n        return (\n            tf.TensorShape(()),\n            tf.TensorShape(()),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape(()),\n            tf.TensorShape(()),\n            tf.TensorShape(()),\n            tf.TensorShape(()),\n            tf.TensorShape(())\n        )\n\n    @property\n    def test_output_type(self) -> Tuple:\n        return (\n            tf.int64,  # unique_id\n            tf.int64,  # example_id\n            tf.int32,  # input_ids\n            tf.int32,  # input_mask\n            tf.int32,  # segment_ids\n            tf.int32,  # token_map\n            tf.int32,  # max_context_map\n        )\n\n    @property\n    def test_output_shape(self) -> Tuple:\n        return (\n            tf.TensorShape(()),\n            tf.TensorShape(()),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n            tf.TensorShape([self.config.max_seq_length]),\n        )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls ../input","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from transformers import BertTokenizer","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dummy_config = Config()\ntokenizer = BertTokenizer(vocab_file='../input/tf2nq-vocab/vocab-nq.txt')\ndummy_preprocessor = NQPreprocessor(dummy_config, tokenizer, mode=Mode.TEST)\nfilepath = Path('../input/tensorflow2-question-answering/simplified-nq-test.jsonl')\ndataloader = NQDataloader(dummy_config, filepath, dummy_preprocessor, mode=Mode.TEST)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"get generator by `generator()` and convert generator into tf.data.Dataset by `to_dataset()`"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ngenerator = dataloader.generator()\n\ncount = 0\nfor _ in generator:\n    count += 1\nprint(count)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ndataset = dataloader.to_dataset(mode=Mode.TEST)\n\ncount = 0\nfor batch in dataset:\n    count += 1\nprint(count)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(batch)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}