{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":12863,"databundleVersionId":788719,"sourceType":"competition"},{"sourceId":7412539,"sourceType":"datasetVersion","datasetId":4311702},{"sourceId":7412605,"sourceType":"datasetVersion","datasetId":4311747},{"sourceId":7412615,"sourceType":"datasetVersion","datasetId":4311755},{"sourceId":7413337,"sourceType":"datasetVersion","datasetId":4312269}],"dockerImageVersionId":30635,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"pip install /kaggle/input/packages/sacremoses-0.1.1-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2024-01-16T12:05:41.854364Z","iopub.execute_input":"2024-01-16T12:05:41.854702Z","iopub.status.idle":"2024-01-16T12:05:54.163369Z","shell.execute_reply.started":"2024-01-16T12:05:41.854673Z","shell.execute_reply":"2024-01-16T12:05:54.162291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pip install /kaggle/input/packages/botocore-1.34.19-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2024-01-16T12:10:59.859273Z","iopub.execute_input":"2024-01-16T12:10:59.859552Z","iopub.status.idle":"2024-01-16T12:11:11.630459Z","shell.execute_reply.started":"2024-01-16T12:10:59.859526Z","shell.execute_reply":"2024-01-16T12:11:11.629327Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pip install /kaggle/input/packages/s3transfer-0.10.0-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2024-01-16T12:12:37.017349Z","iopub.execute_input":"2024-01-16T12:12:37.018205Z","iopub.status.idle":"2024-01-16T12:12:48.723886Z","shell.execute_reply.started":"2024-01-16T12:12:37.018162Z","shell.execute_reply":"2024-01-16T12:12:48.722770Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pip install /kaggle/input/packages/boto3-1.34.19-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2024-01-16T12:13:01.485520Z","iopub.execute_input":"2024-01-16T12:13:01.485870Z","iopub.status.idle":"2024-01-16T12:13:13.282422Z","shell.execute_reply.started":"2024-01-16T12:13:01.485830Z","shell.execute_reply":"2024-01-16T12:13:13.281290Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pip install /kaggle/input/packages/transformers-2.3.0-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2024-01-16T12:13:30.256103Z","iopub.execute_input":"2024-01-16T12:13:30.256469Z","iopub.status.idle":"2024-01-16T12:13:45.467256Z","shell.execute_reply.started":"2024-01-16T12:13:30.256440Z","shell.execute_reply":"2024-01-16T12:13:45.466147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pip show transformers","metadata":{"execution":{"iopub.status.busy":"2024-01-16T12:13:52.482930Z","iopub.execute_input":"2024-01-16T12:13:52.483276Z","iopub.status.idle":"2024-01-16T12:14:03.171168Z","shell.execute_reply.started":"2024-01-16T12:13:52.483247Z","shell.execute_reply":"2024-01-16T12:14:03.170042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nimport argparse\nimport os\nimport random\nimport time\nimport pickle\nimport gc\nimport math\nfrom collections import namedtuple\n\n\nimport numpy as np\nfrom tqdm import tqdm; tqdm.monitor_interval = 0  # noqa\n\nfrom transformers import BertConfig, BertTokenizer, BertModel\nfrom torch import nn\n\n\nimport torch\nimport json\nimport logging\nimport os\nimport collections\nimport pickle\nimport pandas as pd\nfrom tqdm.notebook import tqdm\nimport numpy as np\nfrom transformers import BertModel, BertPreTrainedModel\nimport transformers.tokenization_bert\n\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def whitespace_tokenize(text):\n    text = text.strip()\n    if not text:\n        return []\n    tokens = text.split()\n    return tokens","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\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           \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            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\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\n        special_tokens_count = 3\n        if sep_token_extra:\n\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\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\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=bool)\n            token_type_ids = []\n\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\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            if long_is_impossible:\n                if np.random.rand() > p_keep_impossible:\n                    continue\n\n            tokens.append(cls_token)\n            token_type_ids.append(cls_token_segment_id)\n      \n            tokens += query_tokens\n            token_type_ids += [sequence_a_segment_id] * len(query_tokens)\n        \n            tokens.append(sep_token)\n            token_type_ids.append(sequence_a_segment_id)\n\n            if sep_token_extra:\n                tokens.append(sep_token)\n                token_type_ids.append(sequence_a_segment_id)\n\n            # \n            for i in range(doc_span.length):\n                split_token_index = doc_span.start + i\n             \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    \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=bool)\n            token_type_ids = np.array(token_type_ids, dtype=np.uint8)\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                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        \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    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    tok_text = tok_text.replace(\" ##\", \"\")\n    tok_text = tok_text.replace(\"##\", \"\")\n\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            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    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\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 \n                if crop.token_to_orig_map[long_index] == UNMAPPED:\n                    continue\n            \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        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 \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\n      short_answer = f'{start_token}:{end_token + 1}'\n    df['PredictionString'].append(short_answer)\n\n  \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\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\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\n\nclass BertForNaturalQuestionAnswering(BertPreTrainedModel):\n    def __init__(self, config):\n        super().__init__(config)\n\n        self.num_labels = config.num_labels\n\n        self.bert = BertModel(config)\n        self.qa_outputs = nn.Linear(config.hidden_size, 2)\n        self.long_outputs = nn.Linear(config.hidden_size, 1)\n\n        self.init_weights()\n\n    def forward(self, inputs, **kwargs):\n        outputs = self.bert(\n            #output_hidden_states=True,\n            input_ids=inputs[\"input_ids\"],\n            attention_mask=inputs[\"attention_mask\"],\n            token_type_ids=inputs[\"token_type_ids\"].type(torch.int32)\n        )\n        sequence_output = outputs[0]\n        logits = self.qa_outputs(sequence_output)\n\n        start_logits, end_logits = torch.split(logits, 1, dim=-1)\n        start_logits = start_logits.squeeze(-1)\n        end_logits = end_logits.squeeze(-1)\n        long_logits = self.long_outputs(sequence_output).squeeze(-1)\n\n        return start_logits, end_logits, long_logits\n\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\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\n\nRawResult = namedtuple(\"RawResult\", [\"unique_id\", \"start_logits\", \"end_logits\",\n    \"long_logits\"])\n\n\nclass QADataset(Dataset):\n    def __init__(self, data_dict, is_eval):\n        self.data_dict = data_dict\n        self.is_eval = is_eval\n\n    def __len__(self):\n        return len(self.data_dict[\"input_ids\"])\n\n    def __getitem__(self, index):\n        if self.is_eval:\n            return {\n                'input_ids': self.data_dict[\"input_ids\"][index],\n                'attention_mask': self.data_dict[\"attention_mask\"][index],\n                'token_type_ids': self.data_dict[\"token_type_ids\"][index],\n                'example_index': self.data_dict[\"example_index\"][index]\n            }\n\n        assert False\n\n\ndef submit(args, model, tokenizer):\n    csv_fn = '/kaggle/working/submission.csv'\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    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    eval_dataset_length = len(eval_dataset['input_ids'])\n    padded_length = math.ceil(eval_dataset_length / args.eval_batch_size) * args.eval_batch_size\n    num_pad = padded_length - eval_dataset_length\n    for ti, t in eval_dataset.items():\n        pad_tensor = torch.zeros(size=t[0].shape).reshape(1, -1).type(torch.int32)\n        pad_tensor = pad_tensor.repeat(num_pad, 1)\n        eval_dataset[ti] = torch.concat([t, pad_tensor], 0)\n\n\n    eval_dataset['example_index'] = torch.arange(padded_length).type(torch.int32)\n    eval_ds_aux = QADataset(eval_dataset, True)\n    eval_loader = DataLoader(eval_ds_aux, batch_size=args.eval_batch_size, drop_last=True, shuffle=False)\n    eval_loader = tqdm(eval_loader, desc=\"Batch {}\".format(0))\n\n\n    print(\"***** Running evaluation *****\")\n    print(\"  Num examples = %d\", eval_dataset_length)\n    print(\"  Batch size = %d\", args.eval_batch_size)\n\n    all_results = []\n    tic = time.time()\n    for step, batch in enumerate(eval_loader):\n        batch = {key: value.to(device) for key, value in batch.items()}\n\n        example_indexes = batch['example_index'].tolist()\n\n        with torch.no_grad():\n            outputs = model(batch)\n\n        device_cpu = torch.device(\"cpu\")\n        outputs = [(o.to(device_cpu)).numpy() for o in outputs]\n\n        for i, example_index in enumerate(example_indexes):\n\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 = outputs[0][i]\n            end_logits = outputs[1][i]\n            long_logits = outputs[2][i]\n\n            start_logits = [float(x) for x in start_logits]\n            end_logits = [float(x) for x in end_logits]\n            long_logits = [float(x) for x in long_logits]\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(['/kaggle/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 = '/kaggle/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    all_input_ids = torch.stack([torch.Tensor(c.input_ids) for c in crops], 0).type(torch.int32)\n    all_attention_mask = torch.stack([torch.Tensor(c.attention_mask) for c in crops], 0).type(torch.int32)\n    all_token_type_ids = torch.stack([torch.Tensor(c.token_type_ids) for c in crops], 0).type(torch.int32)\n\n    if evaluate:\n        dataset = {\n            'input_ids': all_input_ids,\n            'attention_mask': all_attention_mask,\n            'token_type_ids': all_token_type_ids\n        }\n\n    else:\n        assert False\n\n    return dataset, crops, entries\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\n    config_basename = os.path.basename(args.model_config)\n\n    do_lower_case = True\n\n\n    args.model_type = args.model_type.lower()\n\n    with open(\"/kaggle/input/bert-tokenizer-config/tokenizer_with_html.pkl\", \"rb\") as fd:\n        tokenizer = pickle.load(fd)\n\n    config = BertConfig.from_pretrained(\"/kaggle/input/bert-model-config\", local_files_only=True)\n    model = BertForNaturalQuestionAnswering(config=config)\n\n    model.load_state_dict(torch.load(\"/kaggle/input/localberttrained/bert_trained_with_amp\"))\n    \n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model.to(device)\n\n    # Evaluate\n    result = submit(args, model, tokenizer)\n    print(\"Result: {}\".format(result))\n\n\nmain()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!head submission.csv","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}