{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install --upgrade ../input/tokenizers-repo/tokenizers-0.0.11-cp36-cp36m-manylinux1_x86_64.whl\n!pip install --upgrade ../input/sacremoses-repo/sacremoses-0.0.35-py3-none-any.whl\n!pip install --upgrade ../input/transformers-repo/transformers-2.3.0-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# coding=utf-8\n# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\" Finetuning the library models for question-answering on SQuAD (DistilBERT, Bert, XLM, XLNet).\"\"\"\nimport argparse\nimport os\nimport random\nimport time\nimport pickle\nimport gc\nimport math\nfrom collections import namedtuple\nimport tensorflow as tf\n\nimport numpy as np\nfrom tqdm import tqdm; tqdm.monitor_interval = 0  # noqa\n\nfrom transformers import BertConfig, BertTokenizer, RobertaConfig, RobertaTokenizer\n\n\n# coding=utf-8\n# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\" Load NQ dataset. \"\"\"\nimport json\nimport logging\nimport os\nimport collections\nimport pickle\nimport pandas as pd\nfrom tqdm.notebook import tqdm\nimport numpy as np\n\nfrom transformers.tokenization_bert import whitespace_tokenize\n\n\nlogger = logging.getLogger(__name__)\n\n\nNQExample = collections.namedtuple(\"NQExample\", [\n    \"qas_id\", \"question_text\", \"doc_tokens\", \"orig_answer_text\",\n    \"start_position\", \"end_position\", \"long_position\",\n    \"short_is_impossible\", \"long_is_impossible\", \"crop_start\"])\n\nCrop = collections.namedtuple(\"Crop\", [\"unique_id\", \"example_index\", \"doc_span_index\",\n    \"tokens\", \"token_to_orig_map\", \"token_is_max_context\",\n    \"input_ids\", \"attention_mask\", \"token_type_ids\",\n    # \"p_mask\",\n    \"paragraph_len\", \"start_position\", \"end_position\", \"long_position\",\n    \"short_is_impossible\", \"long_is_impossible\"])\n\nLongAnswerCandidate = collections.namedtuple('LongAnswerCandidate', [\n    'start_token', 'end_token', 'top_level'])\n\nUNMAPPED = -123\nCLS_INDEX = 0\n\n\ndef get_add_tokens(do_enumerate):\n    tags = ['Dd', 'Dl', 'Dt', 'H1', 'H2', 'H3', 'Li', 'Ol', 'P', 'Table', 'Td', 'Th', 'Tr', 'Ul']\n    opening_tags = [f'<{tag}>' for tag in tags]\n    closing_tags = [f'</{tag}>' for tag in tags]\n    added_tags = opening_tags + closing_tags\n    # See `nq_to_sqaud.py` for special-tokens\n    special_tokens = ['<P>', '<Table>']\n    if do_enumerate:\n        for special_token in special_tokens:\n            for j in range(11):\n              added_tags.append(f'<{special_token[1: -1]}{j}>')\n\n    add_tokens = ['Td_colspan', 'Th_colspan', '``', '\\'\\'', '--']\n    add_tokens = add_tokens + added_tags\n    return add_tokens\n\n\ndef find_closing_tag(tokens, opening_tag):\n    closing_tag = f'</{opening_tag[1: -1]}>'\n    index, stack = -1, []\n    for token_index, token in enumerate(tokens):\n        if token == opening_tag:\n            stack.insert(0, opening_tag)\n        elif token == closing_tag:\n            stack.pop()\n\n        if len(stack) == 0:\n            index = token_index\n            break\n    return index\n\n\ndef read_candidates(candidate_files, do_cache=True):\n    assert isinstance(candidate_files, (tuple, list)), candidate_files\n    for fn in candidate_files:\n        assert os.path.exists(fn), f'Missing file {fn}'\n    cache_fn = 'candidates.pkl'\n\n    candidates = {}\n    if not os.path.exists(cache_fn):\n        for fn in candidate_files:\n            with open(fn) as f:\n                for line in tqdm(f):\n                    entry = json.loads(line)\n                    example_id = str(entry['example_id'])\n                    cnds = entry.pop('long_answer_candidates')\n                    cnds = [LongAnswerCandidate(c['start_token'], c['end_token'],\n                            c['top_level']) for c in cnds]\n                    candidates[example_id] = cnds\n\n        if do_cache:\n            with open(cache_fn, 'wb') as f:\n                pickle.dump(candidates, f)\n    else:\n        print(f'Loading from cache: {cache_fn}')\n        with open(cache_fn, 'rb') as f:\n            candidates = pickle.load(f)\n\n    return candidates\n\n\ndef is_whitespace(c):\n    if c == \" \" or c == \"\\t\" or c == \"\\r\" or c == \"\\n\" or ord(c) == 0x202F:\n        return True\n    return False\n\n\ndef read_nq_examples(input_file_or_data, is_training):\n    \"\"\"Read a NQ json file into a list of NQExample. Refer to `nq_to_squad.py`\n       to convert the `simplified-nq-t*.jsonl` files to NQ json.\"\"\"\n    if isinstance(input_file_or_data, str):\n        with open(input_file_or_data, \"r\", encoding='utf-8') as f:\n            input_data = json.load(f)[\"data\"]\n\n    else:\n        input_data = input_file_or_data\n\n    for entry_index, entry in enumerate(tqdm(input_data, total=len(input_data))):\n        # if entry_index >= 2:\n        #     break\n        assert len(entry[\"paragraphs\"]) == 1\n        paragraph = entry[\"paragraphs\"][0]\n        paragraph_text = paragraph[\"context\"]\n        doc_tokens = []\n        char_to_word_offset = []\n        prev_is_whitespace = True\n        for c in paragraph_text:\n            if is_whitespace(c):\n                prev_is_whitespace = True\n            else:\n                if prev_is_whitespace:\n                    doc_tokens.append(c)\n                else:\n                    doc_tokens[-1] += c\n                prev_is_whitespace = False\n            char_to_word_offset.append(len(doc_tokens) - 1)\n\n        assert len(paragraph[\"qas\"]) == 1\n        qa = paragraph[\"qas\"][0]\n        start_position = None\n        end_position = None\n        long_position = None\n        orig_answer_text = None\n        short_is_impossible = False\n        long_is_impossible = False\n        if is_training:\n            short_is_impossible = qa[\"short_is_impossible\"]\n            short_answers = qa[\"short_answers\"]\n            if len(short_answers) >= 2:\n                # logger.info(f\"Choosing leftmost of \"\n                #     f\"{len(short_answers)} short answer\")\n                short_answers = sorted(short_answers, key=lambda sa: sa[\"answer_start\"])\n                short_answers = short_answers[0: 1]\n\n            if not short_is_impossible:\n                answer = short_answers[0]\n                orig_answer_text = answer[\"text\"]\n                answer_offset = answer[\"answer_start\"]\n                answer_length = len(orig_answer_text)\n                start_position = char_to_word_offset[answer_offset]\n                end_position = char_to_word_offset[\n                    answer_offset + answer_length - 1]\n                # Only add answers where the text can be exactly\n                # recovered from the document. If this CAN'T\n                # happen it's likely due to weird Unicode stuff\n                # so we will just skip the example.\n                #\n                # Note that this means for training mode, every\n                # example is NOT guaranteed to be preserved.\n                actual_text = \" \".join(doc_tokens[start_position:\n                    end_position + 1])\n                cleaned_answer_text = \" \".join(\n                    whitespace_tokenize(orig_answer_text))\n                if actual_text.find(cleaned_answer_text) == -1:\n                    logger.warning(\n                        \"Could not find answer: '%s' vs. '%s'\",\n                        actual_text, cleaned_answer_text)\n                    continue\n            else:\n                start_position = -1\n                end_position = -1\n                orig_answer_text = \"\"\n\n            long_is_impossible = qa[\"long_is_impossible\"]\n            long_answers = qa[\"long_answers\"]\n            if (len(long_answers) != 1) and not long_is_impossible:\n                raise ValueError(f\"For training, each question\"\n                    f\" should have exactly 1 long answer.\")\n\n            if not long_is_impossible:\n                long_answer = long_answers[0]\n                long_answer_offset = long_answer[\"answer_start\"]\n                long_position = char_to_word_offset[long_answer_offset]\n            else:\n                long_position = -1\n\n            # print(f'Q:{question_text}')\n            # print(f'A:{start_position}, {end_position},\n            # {orig_answer_text}')\n            # print(f'R:{doc_tokens[start_position: end_position]}')\n\n            if not short_is_impossible and not long_is_impossible:\n                assert long_position <= start_position\n\n            if not short_is_impossible and long_is_impossible:\n                assert False, f'Invalid pair short, long pair'\n\n        example = NQExample(\n            qas_id=qa[\"id\"],\n            question_text=qa[\"question\"],\n            doc_tokens=doc_tokens,\n            orig_answer_text=orig_answer_text,\n            start_position=start_position,\n            end_position=end_position,\n            long_position=long_position,\n            short_is_impossible=short_is_impossible,\n            long_is_impossible=long_is_impossible,\n            crop_start=qa[\"crop_start\"])\n\n        yield example\n\n\nDocSpan = collections.namedtuple(\"DocSpan\", [\"start\", \"length\"])\n\n\ndef get_spans(doc_stride, max_tokens_for_doc, max_len):\n    doc_spans = []\n    start_offset = 0\n    while start_offset < max_len:\n        length = max_len - start_offset\n        if length > max_tokens_for_doc:\n            length = max_tokens_for_doc\n        doc_spans.append(DocSpan(start=start_offset, length=length))\n        if start_offset + length == max_len:\n            break\n        start_offset += min(length, doc_stride)\n    return doc_spans\n\n\ndef convert_examples_to_crops(examples_gen, tokenizer, max_seq_length,\n                              doc_stride, max_query_length, is_training,\n                              cls_token='[CLS]', sep_token='[SEP]', pad_id=0,\n                              sequence_a_segment_id=0,\n                              sequence_b_segment_id=1,\n                              cls_token_segment_id=0,\n                              pad_token_segment_id=0,\n                              mask_padding_with_zero=True,\n                              p_keep_impossible=None,\n                              sep_token_extra=False):\n    \"\"\"Loads a data file into a list of `InputBatch`s.\"\"\"\n    assert p_keep_impossible is not None, '`p_keep_impossible` is required'\n    unique_id = 1000000000\n    num_short_pos, num_short_neg = 0, 0\n    num_long_pos, num_long_neg = 0, 0\n    sub_token_cache = {}\n    # max_N, max_M = 1024, 1024\n    # f = np.zeros((max_N, max_M), dtype=np.float32)\n\n    crops = []\n    for example_index, example in enumerate(examples_gen):\n        if example_index % 1000 == 0 and example_index > 0:\n            logger.info('Converting %s: short_pos %s short_neg %s'\n                ' long_pos %s long_neg %s',\n                example_index, num_short_pos, num_short_neg,\n                num_long_pos, num_long_neg)\n\n        query_tokens = tokenizer.tokenize(example.question_text)\n        if len(query_tokens) > max_query_length:\n            query_tokens = query_tokens[0:max_query_length]\n\n        # this takes the longest!\n        tok_to_orig_index = []\n        orig_to_tok_index = []\n        all_doc_tokens = []\n\n        for i, token in enumerate(example.doc_tokens):\n            orig_to_tok_index.append(len(all_doc_tokens))\n            sub_tokens = sub_token_cache.get(token)\n            if sub_tokens is None:\n                sub_tokens = tokenizer.tokenize(token)\n                sub_token_cache[token] = sub_tokens\n            tok_to_orig_index.extend([i for _ in range(len(sub_tokens))])\n            all_doc_tokens.extend(sub_tokens)\n\n        tok_start_position = None\n        tok_end_position = None\n        if is_training and example.short_is_impossible:\n            tok_start_position = -1\n            tok_end_position = -1\n\n        if is_training and not example.short_is_impossible:\n            tok_start_position = orig_to_tok_index[example.start_position]\n            if example.end_position < len(example.doc_tokens) - 1:\n                tok_end_position = orig_to_tok_index[\n                    example.end_position + 1] - 1\n            else:\n                tok_end_position = len(all_doc_tokens) - 1\n\n        tok_long_position = None\n        if is_training and example.long_is_impossible:\n            tok_long_position = -1\n\n        if is_training and not example.long_is_impossible:\n            tok_long_position = orig_to_tok_index[example.long_position]\n\n        # For Bert: [CLS] question [SEP] paragraph [SEP]\n        special_tokens_count = 3\n        if sep_token_extra:\n            # For Roberta: <s> question </s> </s> paragraph </s>\n            special_tokens_count += 1\n        max_tokens_for_doc = max_seq_length - len(query_tokens) - special_tokens_count\n        assert max_tokens_for_doc > 0\n        # We can have documents that are longer than the maximum\n        # sequence length. To deal with this we do a sliding window\n        # approach, where we take chunks of the up to our max length\n        # with a stride of `doc_stride`.\n        doc_spans = get_spans(doc_stride, max_tokens_for_doc, len(all_doc_tokens))\n        for doc_span_index, doc_span in enumerate(doc_spans):\n            # Tokens are constructed as: CLS Query SEP Paragraph SEP\n            tokens = []\n            token_to_orig_map = UNMAPPED * np.ones((max_seq_length, ), dtype=np.int32)\n            token_is_max_context = np.zeros((max_seq_length, ), dtype=np.bool)\n            token_type_ids = []\n\n            # p_mask: mask with 1 for token than cannot be in the\n            # answer (0 for token which can be in an answer)\n            # Original TF implem also keep the classification token\n            # (set to 0) (not sure why...)\n            # p_mask = []\n\n            short_is_impossible = example.short_is_impossible\n            start_position = None\n            end_position = None\n            special_tokens_offset = special_tokens_count - 1\n            doc_offset = len(query_tokens) + special_tokens_offset\n            if is_training and not short_is_impossible:\n                doc_start = doc_span.start\n                doc_end = doc_span.start + doc_span.length - 1\n                if not (tok_start_position >= doc_start and tok_end_position <= doc_end):\n                    start_position = 0\n                    end_position = 0\n                    short_is_impossible = True\n                else:\n                    start_position = tok_start_position - doc_start + doc_offset\n                    end_position = tok_end_position - doc_start + doc_offset\n\n            long_is_impossible = example.long_is_impossible\n            long_position = None\n            if is_training and not long_is_impossible:\n                doc_start = doc_span.start\n                doc_end = doc_span.start + doc_span.length - 1\n                # out of span\n                if not (tok_long_position >= doc_start and tok_long_position <= doc_end):\n                    long_position = 0\n                    long_is_impossible = True\n                else:\n                    long_position = tok_long_position - doc_start + doc_offset\n\n            # drop impossible samples\n            if long_is_impossible:\n                if np.random.rand() > p_keep_impossible:\n                    continue\n\n            # CLS token at the beginning\n            tokens.append(cls_token)\n            token_type_ids.append(cls_token_segment_id)\n            # p_mask.append(0)  # can be answer\n\n            # Query\n            tokens += query_tokens\n            token_type_ids += [sequence_a_segment_id] * len(query_tokens)\n            # p_mask += [1] * len(query_tokens)  # can not be answer\n\n            # SEP token\n            tokens.append(sep_token)\n            token_type_ids.append(sequence_a_segment_id)\n            # p_mask.append(1)  # can not be answer\n            if sep_token_extra:\n                tokens.append(sep_token)\n                token_type_ids.append(sequence_a_segment_id)\n                # p_mask.append(1)\n\n            # Paragraph\n            for i in range(doc_span.length):\n                split_token_index = doc_span.start + i\n                # We add `example.crop_start` as the original document\n                # is already shifted\n                token_to_orig_map[len(tokens)] = tok_to_orig_index[\n                    split_token_index] + example.crop_start\n\n                token_is_max_context[len(tokens)] = check_is_max_context(doc_spans,\n                    doc_span_index, split_token_index)\n                tokens.append(all_doc_tokens[split_token_index])\n                token_type_ids.append(sequence_b_segment_id)\n                # p_mask.append(0)  # can be answer\n\n            paragraph_len = doc_span.length\n\n            # SEP token\n            tokens.append(sep_token)\n            token_type_ids.append(sequence_b_segment_id)\n            # p_mask.append(1)  # can not be answer\n\n            input_ids = tokenizer.convert_tokens_to_ids(tokens)\n\n            # The mask has 1 for real tokens and 0 for padding tokens. Only real\n            # tokens are attended to.\n            attention_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)\n\n            # Zero-pad up to the sequence length.\n            while len(input_ids) < max_seq_length:\n                input_ids.append(pad_id)\n                attention_mask.append(0 if mask_padding_with_zero else 1)\n                token_type_ids.append(pad_token_segment_id)\n                # p_mask.append(1)  # can not be answer\n\n            # reduce memory, only input_ids needs more bits\n            input_ids = np.array(input_ids, dtype=np.int32)\n            attention_mask = np.array(attention_mask, dtype=np.bool)\n            token_type_ids = np.array(token_type_ids, dtype=np.uint8)\n            # p_mask = np.array(p_mask, dtype=np.bool)\n\n            if is_training and short_is_impossible:\n                start_position = CLS_INDEX\n                end_position = CLS_INDEX\n\n            if is_training and long_is_impossible:\n                long_position = CLS_INDEX\n\n            if example_index in (0, 10):\n                # too spammy otherwise\n                if doc_span_index in (0, 5):\n                    logger.info(\"*** Example ***\")\n                    logger.info(\"unique_id: %s\" % (unique_id))\n                    logger.info(\"example_index: %s\" % (example_index))\n                    logger.info(\"doc_span_index: %s\" % (doc_span_index))\n                    logger.info(\"tokens: %s\" % \" \".join(tokens))\n                    # logger.info(\"token_to_orig_map: %s\" % \" \".join([\n                    #     \"%d:%d\" % (x, y) for (x, y) in enumerate(token_to_orig_map)]))\n                    # logger.info(\"token_is_max_context: %s\" % \" \".join([\n                    #     \"%d:%s\" % (x, y) for (x, y) in enumerate(token_is_max_context)\n                    # ]))\n                    logger.info(\"input_ids: %s\" % input_ids)\n                    logger.info(\"attention_mask: %s\" % np.uint8(attention_mask))\n                    logger.info(\"token_type_ids: %s\" % token_type_ids)\n                    if is_training and short_is_impossible:\n                        logger.info(\"short impossible example\")\n                    if is_training and long_is_impossible:\n                        logger.info(\"long impossible example\")\n                    if is_training and not short_is_impossible:\n                        answer_text = \" \".join(tokens[start_position: end_position + 1])\n                        logger.info(\"start_position: %d\" % (start_position))\n                        logger.info(\"end_position: %d\" % (end_position))\n                        logger.info(\"answer: %s\" % (answer_text))\n\n            if short_is_impossible:\n                num_short_neg += 1\n            else:\n                num_short_pos += 1\n\n            if long_is_impossible:\n                num_long_neg += 1\n            else:\n                num_long_pos += 1\n\n            crop = Crop(\n                unique_id=unique_id,\n                example_index=example_index,\n                doc_span_index=doc_span_index,\n                tokens=tokens,\n                token_to_orig_map=token_to_orig_map,\n                token_is_max_context=token_is_max_context,\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                token_type_ids=token_type_ids,\n                # p_mask=p_mask,\n                paragraph_len=paragraph_len,\n                start_position=start_position,\n                end_position=end_position,\n                long_position=long_position,\n                short_is_impossible=short_is_impossible,\n                long_is_impossible=long_is_impossible)\n            crops.append(crop)\n            unique_id += 1\n\n    return crops\n\n\ndef check_is_max_context(doc_spans, cur_span_index, position):\n    \"\"\"Check if this is the 'max context' doc span for the token.\"\"\"\n\n    # Because of the sliding window approach taken to scoring documents, a single\n    # token can appear in multiple documents. E.g.\n    #  Doc: the man went to the store and bought a gallon of milk\n    #  Span A: the man went to the\n    #  Span B: to the store and bought\n    #  Span C: and bought a gallon of\n    #  ...\n    #\n    # Now the word 'bought' will have two scores from spans B and C. We only\n    # want to consider the score with \"maximum context\", which we define as\n    # the *minimum* of its left and right context (the *sum* of left and\n    # right context will always be the same, of course).\n    #\n    # In the example the maximum context for 'bought' would be span C since\n    # it has 1 left context and 3 right context, while span B has 4 left context\n    # and 0 right context.\n    best_score = None\n    best_span_index = None\n    for (span_index, doc_span) in enumerate(doc_spans):\n        end = doc_span.start + doc_span.length - 1\n        if position < doc_span.start:\n            continue\n        if position > end:\n            continue\n        num_left_context = position - doc_span.start\n        num_right_context = end - position\n        score = min(num_left_context, num_right_context) + 0.01 * doc_span.length\n        if best_score is None or score > best_score:\n            best_score = score\n            best_span_index = span_index\n\n    return cur_span_index == best_span_index\n\n\nPrelimPrediction = collections.namedtuple(\"PrelimPrediction\",\n    [\"crop_index\", \"start_index\", \"end_index\", \"start_logit\", \"end_logit\"])\n\nNbestPrediction = collections.namedtuple(\"NbestPrediction\", [\n    \"text\", \"start_logit\", \"end_logit\",\n    \"start_index\", \"end_index\",\n    \"orig_doc_start\", \"orig_doc_end\", \"crop_index\"])\n\n\ndef clean_text(tok_text):\n    # De-tokenize WordPieces that have been split off.\n    tok_text = tok_text.replace(\" ##\", \"\")\n    tok_text = tok_text.replace(\"##\", \"\")\n\n    # Clean whitespace\n    tok_text = tok_text.strip()\n    tok_text = \" \".join(tok_text.split())\n    return tok_text\n\n\ndef get_nbest(prelim_predictions, crops, example, n_best_size):\n    seen, nbest = set(), []\n    for pred in prelim_predictions:\n        if len(nbest) >= n_best_size:\n            break\n        crop = crops[pred.crop_index]\n        orig_doc_start, orig_doc_end = -1, -1\n        # non-null\n        orig_doc_start, orig_doc_end = -1, -1\n        if pred.start_index > 0:\n            # Long answer has no end_index. We still generate some text to check\n            if pred.end_index == -1:\n                tok_tokens = crop.tokens[pred.start_index: pred.start_index + 11]\n            else:\n                tok_tokens = crop.tokens[pred.start_index: pred.end_index + 1]\n            tok_text = \" \".join(tok_tokens)\n            tok_text = clean_text(tok_text)\n\n            orig_doc_start = int(crop.token_to_orig_map[pred.start_index])\n            if pred.end_index == -1:\n                orig_doc_end = orig_doc_start + 10\n            else:\n                orig_doc_end = int(crop.token_to_orig_map[pred.end_index])\n\n            final_text = tok_text\n            if final_text in seen:\n                continue\n\n        else:\n            final_text = \"\"\n\n        seen.add(final_text)\n        nbest.append(NbestPrediction(\n            text=final_text,\n            start_logit=pred.start_logit, end_logit=pred.end_logit,\n            start_index=pred.start_index, end_index=pred.end_index,\n            orig_doc_start=orig_doc_start, orig_doc_end=orig_doc_end,\n            crop_index=pred.crop_index))\n\n    # Degenerate case. I never saw this happen.\n    if len(nbest) in (0, 1):\n        nbest.insert(0, NbestPrediction(text=\"empty\",\n            start_logit=0.0, end_logit=0.0,\n            start_index=-1, end_index=-1,\n            orig_doc_start=-1, orig_doc_end=-1,\n            crop_index=UNMAPPED))\n\n    assert len(nbest) >= 1\n    return nbest\n\n\ndef write_predictions(examples_gen, all_crops, all_results, n_best_size,\n                      max_answer_length, output_prediction_file,\n                      output_nbest_file, output_null_log_odds_file, verbose_logging,\n                      short_null_score_diff, long_null_score_diff):\n    \"\"\"Write final predictions to the json file and log-odds of null if needed.\"\"\"\n    logger.info(\"Writing predictions to: %s\" % output_prediction_file)\n    logger.info(\"Writing nbest to: %s\" % output_nbest_file)\n\n    # create indexes\n    example_index_to_crops = collections.defaultdict(list)\n    for crop in all_crops:\n        example_index_to_crops[crop.example_index].append(crop)\n    unique_id_to_result = {result.unique_id: result for result in all_results}\n\n    all_predictions = collections.OrderedDict()\n    all_nbest_json = collections.OrderedDict()\n    scores_diff_json = collections.OrderedDict()\n    short_num_empty, long_num_empty = 0, 0\n    for example_index, example in enumerate(examples_gen):\n        if example_index % 1000 == 0 and example_index > 0:\n            logger.info(f'[{example_index}]: {short_num_empty} short and {long_num_empty} long empty')\n\n        crops = example_index_to_crops[example_index]\n        short_prelim_predictions, long_prelim_predictions = [], []\n        for crop_index, crop in enumerate(crops):\n            assert crop.unique_id in unique_id_to_result, f\"{crop.unique_id}\"\n            result = unique_id_to_result[crop.unique_id]\n            # get the `n_best_size` largest indexes\n            # https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array#23734295\n            start_indexes = np.argpartition(result.start_logits, -n_best_size)[-n_best_size:]\n            start_indexes = [int(x) for x in start_indexes]\n            end_indexes = np.argpartition(result.end_logits, -n_best_size)[-n_best_size:]\n            end_indexes = [int(x) for x in end_indexes]\n\n            # create short answers\n            for start_index in start_indexes:\n                if start_index >= len(crop.tokens):\n                    continue\n                # this skips [CLS] i.e. null prediction\n                if crop.token_to_orig_map[start_index] == UNMAPPED:\n                    continue\n                if not crop.token_is_max_context[start_index]:\n                    continue\n\n                for end_index in end_indexes:\n                    if end_index >= len(crop.tokens):\n                        continue\n                    if crop.token_to_orig_map[end_index] == UNMAPPED:\n                        continue\n                    if end_index < start_index:\n                        continue\n                    length = end_index - start_index + 1\n                    if length > max_answer_length:\n                        continue\n\n                    short_prelim_predictions.append(PrelimPrediction(\n                        crop_index=crop_index,\n                        start_index=start_index,\n                        end_index=end_index,\n                        start_logit=result.start_logits[start_index],\n                        end_logit=result.end_logits[end_index]))\n\n            long_indexes = np.argpartition(result.long_logits, -n_best_size)[-n_best_size:].tolist()\n            for long_index in long_indexes:\n                if long_index >= len(crop.tokens):\n                    continue\n                # this skips [CLS] i.e. null prediction\n                if crop.token_to_orig_map[long_index] == UNMAPPED:\n                    continue\n                # TODO(see--): Is this needed?\n                # -> Yep helps both short and long by about 0.1\n                if not crop.token_is_max_context[long_index]:\n                    continue\n                long_prelim_predictions.append(PrelimPrediction(\n                    crop_index=crop_index,\n                    start_index=long_index, end_index=-1,\n                    start_logit=result.long_logits[long_index],\n                    end_logit=result.long_logits[long_index]))\n\n        short_prelim_predictions = sorted(short_prelim_predictions,\n            key=lambda x: x.start_logit + x.end_logit, reverse=True)\n\n        short_nbest = get_nbest(short_prelim_predictions, crops,\n            example, n_best_size)\n\n        short_best_non_null = None\n        for entry in short_nbest:\n            if short_best_non_null is None:\n                if entry.text != \"\":\n                    short_best_non_null = entry\n\n        long_prelim_predictions = sorted(long_prelim_predictions,\n            key=lambda x: x.start_logit, reverse=True)\n\n        long_nbest = get_nbest(long_prelim_predictions, crops,\n            example, n_best_size)\n\n        long_best_non_null = None\n        for entry in long_nbest:\n            if long_best_non_null is None:\n                if entry.text != \"\":\n                    long_best_non_null = entry\n\n        nbest_json = {'short': [], 'long': []}\n        for kk, entries in [('short', short_nbest), ('long', long_nbest)]:\n            for i, entry in enumerate(entries):\n                output = {}\n                output[\"text\"] = entry.text\n                output[\"start_logit\"] = entry.start_logit\n                output[\"end_logit\"] = entry.end_logit\n                output[\"start_index\"] = entry.start_index\n                output[\"end_index\"] = entry.end_index\n                output[\"orig_doc_start\"] = entry.orig_doc_start\n                output[\"orig_doc_end\"] = entry.orig_doc_end\n                nbest_json[kk].append(output)\n\n        assert len(nbest_json['short']) >= 1\n        assert len(nbest_json['long']) >= 1\n\n        # We use the [CLS] score of the crop that has the maximum positive score\n        # long_score_diff = min_long_score_null - long_best_non_null.start_logit\n        # Predict \"\" if null score - the score of best non-null > threshold\n        try:\n            crop_unique_id = crops[short_best_non_null.crop_index].unique_id\n            start_score_null = unique_id_to_result[crop_unique_id].start_logits[CLS_INDEX]\n            end_score_null = unique_id_to_result[crop_unique_id].end_logits[CLS_INDEX]\n            short_score_null = start_score_null + end_score_null\n            short_score_diff = short_score_null - (short_best_non_null.start_logit +\n                short_best_non_null.end_logit)\n\n            if short_score_diff > short_null_score_diff:\n                final_pred = (\"\", -1, -1)\n                short_num_empty += 1\n            else:\n                final_pred = (short_best_non_null.text, short_best_non_null.orig_doc_start,\n                    short_best_non_null.orig_doc_end)\n        except Exception as e:\n            print(e)\n            final_pred = (\"\", -1, -1)\n            short_num_empty += 1\n\n        try:\n            long_score_null = unique_id_to_result[crops[\n                long_best_non_null.crop_index].unique_id].long_logits[CLS_INDEX]\n            long_score_diff = long_score_null - long_best_non_null.start_logit\n            scores_diff_json[example.qas_id] = {'short_score_diff': short_score_diff,\n                'long_score_diff': long_score_diff}\n\n            if long_score_diff > long_null_score_diff:\n                final_pred += (\"\", -1)\n                long_num_empty += 1\n                # print(f\"LONG EMPTY: {round(long_score_null, 2)} vs \"\n                #     f\"{round(long_best_non_null.start_logit, 2)} (th {long_null_score_diff})\")\n\n            else:\n                final_pred += (long_best_non_null.text, long_best_non_null.orig_doc_start)\n\n        except Exception as e:\n            print(e)\n            final_pred += (\"\", -1)\n            long_num_empty += 1\n\n        all_predictions[example.qas_id] = final_pred\n        all_nbest_json[example.qas_id] = nbest_json\n\n    if output_prediction_file is not None:\n        with open(output_prediction_file, \"w\") as writer:\n            writer.write(json.dumps(all_predictions, indent=2))\n\n    if output_nbest_file is not None:\n        with open(output_nbest_file, \"w\") as writer:\n            writer.write(json.dumps(all_nbest_json, indent=2))\n\n    if output_null_log_odds_file is not None:\n        with open(output_null_log_odds_file, \"w\") as writer:\n            writer.write(json.dumps(scores_diff_json, indent=2))\n\n    logger.info(f'{short_num_empty} short and {long_num_empty} long empty of'\n        f' {example_index}')\n    return all_predictions\n\n\ndef convert_preds_to_df(preds, candidates):\n  num_found_long, num_searched_long = 0, 0\n  df = {'example_id': [], 'PredictionString': []}\n  for example_id, pred in preds.items():\n    short_text, start_token, end_token, long_text, long_token = pred\n    df['example_id'].append(example_id + '_short')\n    short_answer = ''\n    if start_token != -1:\n      # +1 is required to make the token inclusive\n      short_answer = f'{start_token}:{end_token + 1}'\n    df['PredictionString'].append(short_answer)\n\n    # print(entry['document_text'].split(' ')[start_token: end_token + 1])\n    # find the long answer\n    long_answer = ''\n    found_long = False\n    min_dist = 1_000_000\n    if long_token != -1:\n      num_searched_long += 1\n      for candidate in candidates[example_id]:\n        cstart, cend = candidate.start_token, candidate.end_token\n        dist = abs(cstart - long_token)\n        if dist < min_dist:\n          min_dist = dist\n        if long_token == cstart:\n          long_answer = f'{cstart}:{cend}'\n          found_long = True\n          break\n\n      if found_long:\n        num_found_long += 1\n      else:\n        logger.info(f\"Not found: {min_dist}\")\n\n    df['example_id'].append(example_id + '_long')\n    df['PredictionString'].append(long_answer)\n\n  df = pd.DataFrame(df)\n  print(f'Found {num_found_long} of {num_searched_long} (total {len(preds)})')\n  return df\nimport tensorflow as tf\nfrom tensorflow.keras import layers as L\n\nfrom transformers import TFBertMainLayer, TFBertPreTrainedModel, TFRobertaMainLayer, TFRobertaPreTrainedModel\nfrom transformers.modeling_tf_utils import get_initializer\n\n\nclass TFBertForNaturalQuestionAnswering(TFBertPreTrainedModel):\n    def __init__(self, config, *inputs, **kwargs):\n        super().__init__(config, *inputs, **kwargs)\n        self.num_labels = config.num_labels\n\n        self.bert = TFBertMainLayer(config, name='bert')\n        self.initializer = get_initializer(config.initializer_range)\n        self.qa_outputs = L.Dense(config.num_labels,\n            kernel_initializer=self.initializer, name='qa_outputs')\n        self.long_outputs = L.Dense(1, kernel_initializer=self.initializer,\n            name='long_outputs')\n\n    def call(self, inputs, **kwargs):\n        outputs = self.bert(inputs, **kwargs)\n        sequence_output = outputs[0]\n        logits = self.qa_outputs(sequence_output)\n        start_logits, end_logits = tf.split(logits, 2, axis=-1)\n        start_logits = tf.squeeze(start_logits, -1)\n        end_logits = tf.squeeze(end_logits, -1)\n        long_logits = tf.squeeze(self.long_outputs(sequence_output), -1)\n        return start_logits, end_logits, long_logits\n\n\nclass TFRobertaForNaturalQuestionAnswering(TFRobertaPreTrainedModel):\n    def __init__(self, config, *inputs, **kwargs):\n        super().__init__(config, *inputs, **kwargs)\n        self.num_labels = config.num_labels\n\n        self.roberta = TFRobertaMainLayer(config, name='roberta')\n        self.initializer = get_initializer(config.initializer_range)\n        self.qa_outputs = L.Dense(config.num_labels,\n            kernel_initializer=self.initializer, name='qa_outputs')\n        self.long_outputs = L.Dense(1, kernel_initializer=self.initializer,\n            name='long_outputs')\n\n    def call(self, inputs, **kwargs):\n        outputs = self.roberta(inputs, **kwargs)\n        sequence_output = outputs[0]\n        logits = self.qa_outputs(sequence_output)\n        start_logits, end_logits = tf.split(logits, 2, axis=-1)\n        start_logits = tf.squeeze(start_logits, -1)\n        end_logits = tf.squeeze(end_logits, -1)\n        long_logits = tf.squeeze(self.long_outputs(sequence_output), -1)\n        return start_logits, end_logits, long_logits\n\n\nimport json\nimport numpy as np\nimport argparse\nimport pandas as pd\nfrom tqdm.notebook import tqdm\n\n\ndef enumerate_tags(text_split):\n  \"\"\"Reproduce the preprocessing from:\n  A BERT Baseline for the Natural Questions (https://arxiv.org/pdf/1901.08634.pdf)\n\n  We introduce special markup tokens in the doc-ument  to  give  the  model\n  a  notion  of  which  partof the document it is reading.  The special\n  tokenswe introduced are of the form “[Paragraph=N]”,“[Table=N]”, and “[List=N]”\n  at the beginning ofthe N-th paragraph,  list and table respectively\n  inthe document. This decision was based on the ob-servation that the first\n  few paragraphs and tables inthe document are much more likely than the rest\n  ofthe document to contain the annotated answer andso the model could benefit\n  from knowing whetherit is processing one of these passages.\n\n  We deviate as follows: Tokens are only created for the first 10 times. All other\n  tokens are the same. We only add `special_tokens`. These two are added as they\n  make 72.9% + 19.0% = 91.9% of long answers.\n  (https://github.com/google-research-datasets/natural-questions)\n  \"\"\"\n  special_tokens = ['<P>', '<Table>']\n  special_token_counts = [0 for _ in range(len(special_tokens))]\n  for index, token in enumerate(text_split):\n    for special_token_index, special_token in enumerate(special_tokens):\n      if token == special_token:\n        cnt = special_token_counts[special_token_index]\n        if cnt <= 10:\n          text_split[index] = f'<{special_token[1: -1]}{cnt}>'\n        special_token_counts[special_token_index] = cnt + 1\n\n  return text_split\n\n\ndef convert_nq_to_squad(args=None):\n  np.random.seed(123)\n  if args is None:\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--fn', type=str, default='simplified-nq-train.jsonl')\n    parser.add_argument('--version', type=str, default='v1.0.2')\n    parser.add_argument('--prefix', type=str, default='nq')\n    parser.add_argument('--p_val', type=float, default=0.1)\n    parser.add_argument('--crop_len', type=int, default=2_500)\n    parser.add_argument('--num_samples', type=int, default=1_000_000)\n    parser.add_argument('--val_ids', type=str, default='val_ids.csv')\n    parser.add_argument('--do_enumerate', action='store_true')\n    parser.add_argument('--do_not_dump', action='store_true')\n    parser.add_argument('--num_max_tokens', type=int, default=400_000)\n    args = parser.parse_args()\n\n  is_train = 'train' in args.fn\n  if is_train:\n    train_fn = f'{args.prefix}-train-{args.version}.json'\n    val_fn = f'{args.prefix}-val-{args.version}.json'\n    print(f'Converting {args.fn} to {train_fn} & {val_fn} ... ')\n  else:\n    test_fn = f'{args.prefix}-test-{args.version}.json'\n    print(f'Converting {args.fn} to {test_fn} ... ')\n\n  if args.val_ids:\n    val_ids = set(str(x) for x in pd.read_csv(args.val_ids)['val_ids'].values)\n  else:\n    val_ids = set()\n\n  entries = []\n  smooth = 0.999\n  total_split_len, long_split_len = 0., 0.\n  long_end = 0.\n  num_very_long, num_yes_no, num_short_dropped, num_trimmed = 0, 0, 0, 0\n  num_short_possible, num_long_possible = 0, 0\n  max_end_token = -1\n  orig_data = {}\n  with open(args.fn) as f:\n    progress = tqdm(f, total=args.num_samples)\n    entry = {}\n    for kk, line in enumerate(progress):\n      if kk >= args.num_samples:\n        break\n\n      data = json.loads(line)\n      data_cpy = data.copy()\n      example_id = str(data_cpy.pop('example_id'))\n      data_cpy['document_text'] = ''\n      orig_data[example_id] = data_cpy\n      url = 'MISSING' if not is_train else data['document_url']\n      # progress.write(f'############ {url} ###############')\n      document_text = data['document_text']\n      document_text_split = document_text.split(' ')\n      # trim super long\n      if len(document_text_split) > args.num_max_tokens:\n        num_trimmed += 1\n        document_text_split = document_text_split[:args.num_max_tokens]\n\n      if args.do_enumerate:\n        document_text_split = enumerate_tags(document_text_split)\n      question = data['question_text']  # + '?'\n      annotations = [None] if not is_train else data['annotations']\n      assert len(annotations) == 1, annotations\n      # User str keys!\n      example_id = str(data['example_id'])\n      candidates = data['long_answer_candidates']\n      if not is_train:\n        qa = {'question': question, 'id': example_id, 'crop_start': 0}\n        context = ' '.join(document_text_split)\n\n      else:\n        long_answer = annotations[0]['long_answer']\n        long_answer_len = long_answer['end_token'] - long_answer['start_token']\n        total_split_len = smooth * total_split_len + (1. - smooth) * len(\n            document_text_split)\n        long_split_len = smooth * long_split_len + (1. - smooth) * \\\n            long_answer_len\n        if long_answer['end_token'] > 0:\n          long_end = smooth * long_end + (1. - smooth) * long_answer['end_token']\n\n        if long_answer['end_token'] > max_end_token:\n          max_end_token = long_answer['end_token']\n\n        progress.set_postfix({'ltotal': int(total_split_len),\n            'llong': int(long_split_len), 'long_end': round(long_end, 2)})\n\n        short_answers = annotations[0]['short_answers']\n        yes_no_answer = annotations[0]['yes_no_answer']\n        if yes_no_answer != 'NONE':\n          # progress.write(f'Skipping yes-no: {yes_no_answer}')\n          num_yes_no += 1\n          continue\n\n        # print(f'Q: {question}')\n        # print(f'L: {long_answer_str}')\n        long_is_impossible = long_answer['start_token'] == -1\n        if long_is_impossible:\n          long_answer_candidate = np.random.randint(len(candidates))\n        else:\n          long_answer_candidate = long_answer['candidate_index']\n\n        long_start_token = candidates[long_answer_candidate]['start_token']\n        long_end_token = candidates[long_answer_candidate]['end_token']\n        # generate crop based on tokens. Note that validation samples should\n        # not be cropped as this won't reflect test set performance.\n        if args.crop_len > 0 and example_id not in val_ids:\n          crop_start = long_start_token - np.random.randint(int(args.crop_len * 0.75))\n          if crop_start <= 0:\n            crop_start = 0\n            crop_start_len = -1\n          else:\n            crop_start_len = len(' '.join(document_text_split[:crop_start]))\n\n          crop_end = crop_start + args.crop_len\n        else:\n          crop_start = 0\n          crop_start_len = -1\n          crop_end = 10_000_000\n\n        is_very_long = False\n        if long_end_token > crop_end:\n          num_very_long += 1\n          is_very_long = True\n          # progress.write(f'{num_very_long}: Skipping very long answer {long_end_token}, {crop_end}')\n          # continue\n\n        document_text_crop_split = document_text_split[crop_start: crop_end]\n        context = ' '.join(document_text_crop_split)\n        # create long answer\n        long_answers_ = []\n        if not long_is_impossible:\n          long_answer_pre_split = document_text_split[:long_answer[\n              'start_token']]\n          long_answer_start = len(' '.join(long_answer_pre_split)) - \\\n              crop_start_len\n          long_answer_split = document_text_split[long_answer['start_token']:\n              long_answer['end_token']]\n          long_answer_text = ' '.join(long_answer_split)\n          if not is_very_long:\n            assert context[long_answer_start: long_answer_start + len(\n                long_answer_text)] == long_answer_text, long_answer_text\n          long_answers_ = [{'text': long_answer_text,\n              'answer_start': long_answer_start}]\n\n        # create short answers\n        short_is_impossible = len(short_answers) == 0\n        short_answers_ = []\n        if not short_is_impossible:\n          for short_answer in short_answers:\n            short_start_token = short_answer['start_token']\n            short_end_token = short_answer['end_token']\n            if short_start_token >= crop_start + args.crop_len:\n              num_short_dropped += 1\n              continue\n            short_answers_pre_split = document_text_split[:short_start_token]\n            short_answer_start = len(' '.join(short_answers_pre_split)) - \\\n                crop_start_len\n            short_answer_split = document_text_split[short_start_token: short_end_token]\n            short_answer_text = ' '.join(short_answer_split)\n            assert short_answer_text != ''\n\n            # this happens if we crop and parts of the short answer overflow\n            short_from_context = context[short_answer_start: short_answer_start + len(short_answer_text)]\n            if short_from_context != short_answer_text:\n              print(f'short diff: {short_from_context} vs {short_answer_text}')\n            short_answers_.append({'text': short_from_context,\n                'answer_start': short_answer_start})\n\n        if len(short_answers_) == 0:\n          short_is_impossible = True\n\n        if not short_is_impossible:\n          num_short_possible += 1\n        if not long_is_impossible:\n          num_long_possible += 1\n\n        qa = {'question': question,\n            'short_answers': short_answers_, 'long_answers': long_answers_,\n            'id': example_id, 'short_is_impossible': short_is_impossible,\n            'long_is_impossible': long_is_impossible,\n            'crop_start': crop_start}\n\n      paragraph = {'qas': [qa], 'context': context}\n      entry = {'title': url, 'paragraphs': [paragraph]}\n      entries.append(entry)\n\n  progress.write('  ------------ STATS ------------------')\n  progress.write(f'  Found {num_yes_no} yes/no, {num_very_long} very long'\n      f' and {num_short_dropped} short of {kk} and trimmed {num_trimmed}')\n  progress.write(f'  #short {num_short_possible} #long {num_long_possible}'\n      f' of {len(entries)}')\n  \n  # shuffle to test remaining code\n  np.random.shuffle(entries)\n\n  if is_train:\n    train_entries, val_entries = [], []\n    for entry in entries:\n      if entry['paragraphs'][0]['qas'][0]['id'] not in val_ids:\n        train_entries.append(entry)\n      else:\n        val_entries.append(entry)\n\n    for out_fn, entries in [(train_fn, train_entries), (val_fn, val_entries)]:\n      if not args.do_not_dump:\n        with open(out_fn, 'w') as f:\n          json.dump({'version': args.version, 'data': entries}, f)\n        progress.write(f'Wrote {len(entries)} entries to {out_fn}')\n\n      # save val in competition csv format\n      if 'val' in out_fn:\n        val_example_ids, val_strs = [], []\n        for entry in entries:\n          example_id = entry['paragraphs'][0]['qas'][0]['id']\n          short_answers = orig_data[example_id]['annotations'][0][\n              'short_answers']\n          sa_str = ''\n          for si, sa in enumerate(short_answers):\n            sa_str += f'{sa[\"start_token\"]}:{sa[\"end_token\"]}'\n            if si < len(short_answers) - 1:\n              sa_str += ' '\n          val_example_ids.append(example_id + '_short')\n          val_strs.append(sa_str)\n\n          la = orig_data[example_id]['annotations'][0][\n              'long_answer']\n          la_str = ''\n          if la['start_token'] > 0:\n            la_str += f'{la[\"start_token\"]}:{la[\"end_token\"]}'\n          val_example_ids.append(example_id + '_long')\n          val_strs.append(la_str)\n\n        val_df = pd.DataFrame({'example_id': val_example_ids,\n            'PredictionString': val_strs})\n        val_csv_fn = f'{args.prefix}-val-{args.version}.csv'\n        val_df.to_csv(val_csv_fn, index=False, columns=['example_id',\n            'PredictionString'])\n        print(f'Wrote csv to {val_csv_fn}')\n\n  else:\n    if not args.do_not_dump:\n      with open(test_fn, 'w') as f:\n        json.dump({'version': args.version, 'data': entries}, f)\n      progress.write(f'Wrote to {test_fn}')\n\n  if args.val_ids:\n    print(f'Using val ids from: {args.val_ids}')\n  return entries\n\n\n\nALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys())\n                  for conf in (BertConfig, )), ())\n\nMODEL_CLASSES = {\n    'bert': (BertConfig, TFBertForNaturalQuestionAnswering, BertTokenizer),\n    'roberta': (RobertaConfig, TFRobertaForNaturalQuestionAnswering, RobertaTokenizer),\n}\n\nRawResult = namedtuple(\"RawResult\", [\"unique_id\", \"start_logits\", \"end_logits\",\n    \"long_logits\"])\n\n\ndef set_seed(args):\n    random.seed(args.seed)\n    np.random.seed(args.seed)\n    tf.random.set_seed(args.seed)\n\n\ndef submit(args, model, tokenizer):\n    csv_fn = 'submission.csv'\n    # all_input_ids, all_attention_mask, all_token_type_ids, all_p_mask\n    eval_dataset, crops, entries = load_and_cache_crops(args, tokenizer, evaluate=True)\n    args.eval_batch_size = args.per_tpu_eval_batch_size\n\n    # pad dataset to multiple of `args.eval_batch_size`\n    eval_dataset_length = len(eval_dataset[0])\n    padded_length = math.ceil(eval_dataset_length / args.eval_batch_size) * args.eval_batch_size\n    num_pad = padded_length - eval_dataset[0].shape[0]\n    for ti, t in enumerate(eval_dataset):\n        pad_tensor = tf.expand_dims(tf.zeros_like(t[0]), 0)\n        pad_tensor = tf.repeat(pad_tensor, num_pad, 0)\n        eval_dataset[ti] = tf.concat([t, pad_tensor], 0)\n\n    # create eval dataset\n    eval_ds = tf.data.Dataset.from_tensor_slices({\n        'input_ids': tf.constant(eval_dataset[0]),\n        'attention_mask': tf.constant(eval_dataset[1]),\n        'token_type_ids': tf.constant(eval_dataset[2]),\n        'example_index': tf.range(padded_length, dtype=tf.int32)\n\n    })\n    eval_ds = eval_ds.batch(batch_size=args.eval_batch_size, drop_remainder=True)\n    # eval_ds = eval_ds.prefetch(tf.data.experimental.AUTOTUNE)\n    # eval_ds = strategy.experimental_distribute_dataset(eval_ds)\n\n    # eval\n    print(\"***** Running evaluation *****\")\n    print(\"  Num examples = %d\", eval_dataset_length)\n    print(\"  Batch size = %d\", args.eval_batch_size)\n\n    @tf.function\n    def predict_step(batch):\n        outputs = model(batch, training=False)\n        return outputs\n\n    all_results = []\n    tic = time.time()\n    for batch_ind, batch in tqdm(enumerate(eval_ds), total=padded_length // args.per_tpu_eval_batch_size):\n        # if batch_ind > 2:\n        #     break\n        example_indexes = batch['example_index']\n        # outputs = strategy.experimental_run_v2(predict_step, args=(batch, ))\n        outputs = predict_step(batch)\n        batched_start_logits = outputs[0].numpy()\n        batched_end_logits = outputs[1].numpy()\n        batched_long_logits = outputs[2].numpy()\n        for i, example_index in enumerate(example_indexes):\n            # filter out padded samples\n            if example_index >= eval_dataset_length:\n                continue\n\n            eval_crop = crops[example_index]\n            unique_id = int(eval_crop.unique_id)\n            start_logits = batched_start_logits[i].tolist()\n            end_logits = batched_end_logits[i].tolist()\n            long_logits = batched_long_logits[i].tolist()\n\n            result = RawResult(unique_id=unique_id,\n                               start_logits=start_logits,\n                               end_logits=end_logits,\n                               long_logits=long_logits)\n            all_results.append(result)\n\n    eval_time = time.time() - tic\n    print(\"  Evaluation done in total %f secs (%f sec per example)\",\n        eval_time, eval_time / padded_length)\n    examples_gen = read_nq_examples(entries, is_training=False)\n    preds = write_predictions(examples_gen, crops, all_results, args.n_best_size,\n                              args.max_answer_length,\n                              None, None, None,\n                              args.verbose_logging,\n                              args.short_null_score_diff_threshold, args.long_null_score_diff_threshold)\n    del crops, all_results\n    gc.collect()\n    candidates = read_candidates(['../input/tensorflow2-question-answering/simplified-nq-test.jsonl'], do_cache=False)\n    sub = convert_preds_to_df(preds, candidates).sort_values('example_id')\n    sub.to_csv(csv_fn, index=False, columns=['example_id', 'PredictionString'])\n    print(f'***** Wrote submission to {csv_fn} *****')\n    result = {}\n    return result\n\n\ndef get_convert_args():\n    convert_args = argparse.Namespace()\n    convert_args.fn = '../input/tensorflow2-question-answering/simplified-nq-test.jsonl'\n    convert_args.version = 'v0.0.1'\n    convert_args.prefix = 'nq'\n    convert_args.num_samples = 1_000_000\n    convert_args.val_ids = None\n    convert_args.do_enumerate = False\n    convert_args.do_not_dump = True\n    convert_args.num_max_tokens = 400_000\n    return convert_args\n\n\ndef load_and_cache_crops(args, tokenizer, evaluate=False):\n    # Load data crops from cache or dataset file\n    do_cache = False\n    cached_crops_fn = 'cached_{}_{}.pkl'.format('test', str(args.max_seq_length))\n    if os.path.exists(cached_crops_fn) and do_cache:\n        print(\"Loading crops from cached file %s\", cached_crops_fn)\n        with open(cached_crops_fn, \"rb\") as f:\n            crops = pickle.load(f)\n    else:\n        entries = convert_nq_to_squad(get_convert_args())\n        examples_gen = read_nq_examples(entries, is_training=not evaluate)\n        crops = convert_examples_to_crops(examples_gen=examples_gen,\n                                          tokenizer=tokenizer,\n                                          max_seq_length=args.max_seq_length,\n                                          doc_stride=args.doc_stride,\n                                          max_query_length=args.max_query_length,\n                                          is_training=not evaluate,\n                                          cls_token_segment_id=0,\n                                          pad_token_segment_id=0,\n                                          p_keep_impossible=args.p_keep_impossible if not evaluate else 1.0)\n        if do_cache:\n            with open(cached_crops_fn, \"wb\") as f:\n                pickle.dump(crops, f)\n\n    # stack\n    all_input_ids = tf.stack([c.input_ids for c in crops], 0)\n    all_attention_mask = tf.stack([c.attention_mask for c in crops], 0)\n    all_token_type_ids = tf.stack([c.token_type_ids for c in crops], 0)\n    # all_p_mask = tf.stack([c.p_mask for c in crops], 0)\n\n    # cast `tf.bool`\n    all_attention_mask = tf.cast(all_attention_mask, tf.int32)\n    # all_p_mask = tf.cast(all_p_mask, tf.int32)\n    # all_token_type_ids = tf.cast(all_token_type_ids, tf.int32)\n    if evaluate:\n        dataset = [all_input_ids, all_attention_mask, all_token_type_ids]\n    else:\n        all_start_positions = tf.convert_to_tensor([f.start_position for f in crops], dtype=tf.int32)\n        all_end_positions = tf.convert_to_tensor([f.end_position for f in crops], dtype=tf.int32)\n        all_long_positions = tf.convert_to_tensor([f.long_position for f in crops], dtype=tf.int32)\n        dataset = [all_input_ids, all_attention_mask, all_token_type_ids,\n            all_start_positions, all_end_positions, all_long_positions]\n\n    return dataset, crops, entries\n\n\ndef main():\n    parser = argparse.ArgumentParser()\n\n    # Required parameters\n    parser.add_argument(\"--model_type\", default=\"bert\", type=str)\n    parser.add_argument(\"--model_config\",\n        default=\"../input/transformers-cache/bert_large_uncased_config.json\", type=str)\n    parser.add_argument(\"--checkpoint_dir\", default=\"../input/nq-bert-uncased-68\", type=str)\n    parser.add_argument(\"--vocab_txt\", default=\"../input/transformers-cache/bert_large_uncased_vocab.txt\", type=str)\n\n    # Other parameters\n    parser.add_argument('--short_null_score_diff_threshold', type=float, default=0.0)\n    parser.add_argument('--long_null_score_diff_threshold', type=float, default=0.0)\n    parser.add_argument(\"--max_seq_length\", default=512, type=int)\n    parser.add_argument(\"--doc_stride\", default=256, type=int)\n    parser.add_argument(\"--max_query_length\", default=64, type=int)\n    parser.add_argument(\"--per_tpu_eval_batch_size\", default=4, type=int)\n    parser.add_argument(\"--n_best_size\", default=10, type=int)\n    parser.add_argument(\"--max_answer_length\", default=30, type=int)\n    parser.add_argument(\"--verbose_logging\", action='store_true')\n    parser.add_argument('--seed', type=int, default=42)\n    parser.add_argument('--p_keep_impossible', type=float,\n                        default=0.1, help=\"The fraction of impossible\"\n                        \" samples to keep.\")\n    parser.add_argument('--do_enumerate', action='store_true')\n\n    args, _ = parser.parse_known_args()\n    assert args.model_type not in ('xlnet', 'xlm'), f'Unsupported model_type: {args.model_type}'\n\n    # Set seed\n    set_seed(args)\n\n    # Set cased / uncased\n    config_basename = os.path.basename(args.model_config)\n    if config_basename.startswith('bert'):\n        do_lower_case = 'uncased' in config_basename\n    elif config_basename.startswith('roberta'):\n        # https://github.com/huggingface/transformers/pull/1386/files\n        do_lower_case = False\n\n    # Set XLA\n    # https://github.com/kamalkraj/ALBERT-TF2.0/blob/8d0cc211361e81a648bf846d8ec84225273db0e4/run_classifer.py#L136\n    tf.config.optimizer.set_jit(True)\n    tf.config.optimizer.set_experimental_options({'pin_to_host_optimization': False})\n\n    print(\"Training / evaluation parameters %s\", args)\n    args.model_type = args.model_type.lower()\n    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]\n    config = config_class.from_json_file(args.model_config)\n    tokenizer = tokenizer_class(args.vocab_txt, do_lower_case=do_lower_case)\n    tags = get_add_tokens(do_enumerate=args.do_enumerate)\n    num_added = tokenizer.add_tokens(tags, offset=1)\n    print(f\"Added {num_added} tokens\")\n    print(\"Evaluate the following checkpoint: %s\", args.checkpoint_dir)\n    weights_fn = os.path.join(args.checkpoint_dir, 'weights.h5')\n    model = model_class(config)\n    model(model.dummy_inputs, training=False)\n    model.load_weights(weights_fn)\n\n    # Evaluate\n    result = submit(args, model, tokenizer)\n    print(\"Result: {}\".format(result))\n\n\nmain()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!head submission.csv","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}