{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.6.6","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":769030,"sourceType":"datasetVersion","datasetId":400777},{"sourceId":8908188,"sourceType":"datasetVersion","datasetId":5356177}],"dockerImageVersionId":29841,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nprint(torch.version.cuda)\nprint(torch.__version__)\ntorch.cuda.is_available()","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:05:07.183417Z","iopub.execute_input":"2024-07-08T18:05:07.18387Z","iopub.status.idle":"2024-07-08T18:05:08.440085Z","shell.execute_reply.started":"2024-07-08T18:05:07.183783Z","shell.execute_reply":"2024-07-08T18:05:08.438947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ../input/apex-master/apex\n!pip install --no-cache-dir transformers","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-08T18:05:08.441797Z","iopub.execute_input":"2024-07-08T18:05:08.44214Z","iopub.status.idle":"2024-07-08T18:10:06.589452Z","shell.execute_reply.started":"2024-07-08T18:05:08.44208Z","shell.execute_reply":"2024-07-08T18:10:06.588547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from collections import defaultdict\nfrom dataclasses import dataclass\nimport functools\nimport gc\nimport itertools\nimport json\nfrom multiprocessing import Pool\nimport os\nfrom pathlib import Path\nimport random\nimport re\nimport shutil\nimport subprocess\nimport time\nfrom typing import Callable, Dict, List, Generator, Tuple\n\nimport numpy as np\nimport pandas as pd\nfrom pandas.io.json._json import JsonReader\nfrom sklearn.preprocessing import LabelEncoder\nfrom tqdm._tqdm_notebook import tqdm_notebook as tqdm\n\nimport torch\nfrom torch import nn, optim\nfrom torch.utils.data import Dataset, Subset, DataLoader\n\nfrom apex import amp\nfrom transformers import BertTokenizer, AdamW, get_linear_schedule_with_warmup, BertModel, BertPreTrainedModel","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-07-08T18:03:55.574045Z","iopub.execute_input":"2024-07-08T18:03:55.574386Z","iopub.status.idle":"2024-07-08T18:03:56.704178Z","shell.execute_reply.started":"2024-07-08T18:03:55.574357Z","shell.execute_reply":"2024-07-08T18:03:56.701479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_DIR = Path('../input/tensorflow2-question-answering/')\nDATA_PATH = DATA_DIR / 'simplified-nq-train.jsonl'\n\nstart_time = time.time()\n\nseed = 1029\nvalid_size = 0\ntrain_size = 307373 - valid_size\n\nchunksize = 1000\nmax_seq_len = 384\nmax_question_len = 64\ndoc_stride = 128\n\nnum_labels = 5\nn_epochs = 1\nlr = 2e-5\nwarmup = 0.05\nbatch_size = 16\naccumulation_steps = 4\n\nbert_model = 'bert-base-uncased'\ndo_lower_case = 'uncased' in bert_model\ndevice = torch.device('cuda')\n\noutput_model_file = 'bert_pytorch.bin'\noutput_optimizer_file = 'bert_pytorch_optimizer.bin'\noutput_amp_file = 'bert_pytorch_amp.bin'\n\nrandom.seed(seed)\nnp.random.seed(seed)\ntorch.manual_seed(seed)\ntorch.cuda.manual_seed(seed)\ntorch.backends.cudnn.deterministic = True","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.704879Z","iopub.status.idle":"2024-07-08T18:03:56.705195Z","shell.execute_reply.started":"2024-07-08T18:03:56.705036Z","shell.execute_reply":"2024-07-08T18:03:56.705049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@dataclass\nclass Example(object):\n    example_id: int\n    candidates: List[Dict]\n    annotations: Dict\n    doc_start: int\n    question_len: int\n    tokenized_to_original_index: List[int]\n    input_ids: List[int]\n    start_position: int\n    end_position: int\n    class_label: str\n\n        \ndef convert_data(\n    line: str,\n    tokenizer: BertTokenizer,\n    max_seq_len: int,\n    max_question_len: int,\n    doc_stride: int\n) -> List[Example]:\n    \"\"\"Convert dictionary data into list of training data.\n\n    Parameters\n    ----------\n    line : str\n        Training data.\n    tokenizer : transformers.BertTokenizer\n        Tokenizer for encoding texts into ids.\n    max_seq_len : int\n        Maximum input sequence length.\n    max_question_len : int\n        Maximum input question length.\n    doc_stride : int\n        When splitting up a long document into chunks, how much stride to take between chunks.\n    \"\"\"\n\n    def _find_short_range(short_answers: List[Dict]) -> Tuple[int, int]:\n        answers = pd.DataFrame(short_answers)\n        start_min = answers['start_token'].min()\n        end_max = answers['end_token'].max()\n        return start_min, end_max\n\n    # model input\n    data = json.loads(line)\n    doc_words = data['document_text'].split()\n    question_tokens = tokenizer.tokenize(data['question_text'])[:max_question_len]\n\n    # tokenized index of i-th original token corresponds to original_to_tokenized_index[i]\n    # if a token in original text is removed, its tokenized index indicates next token\n    original_to_tokenized_index = []\n    tokenized_to_original_index = []\n    all_doc_tokens = []  # tokenized document text\n    for i, word in enumerate(doc_words):\n        original_to_tokenized_index.append(len(all_doc_tokens))\n        if re.match(r'<.+>', word):  # remove paragraph tag\n            continue\n        sub_tokens = tokenizer.tokenize(word)\n        for sub_token in sub_tokens:\n            tokenized_to_original_index.append(i)\n            all_doc_tokens.append(sub_token)\n\n    # model output: (class_label, start_position, end_position)\n    annotations = data['annotations'][0]\n    if annotations['yes_no_answer'] in ['YES', 'NO']:\n        class_label = annotations['yes_no_answer'].lower()\n        start_position = annotations['long_answer']['start_token']\n        end_position = annotations['long_answer']['end_token']\n    elif annotations['short_answers']:\n        class_label = 'short'\n        start_position, end_position = _find_short_range(annotations['short_answers'])\n    elif annotations['long_answer']['candidate_index'] != -1:\n        class_label = 'long'\n        start_position = annotations['long_answer']['start_token']\n        end_position = annotations['long_answer']['end_token']\n    else:\n        class_label = 'unknown'\n        start_position = -1\n        end_position = -1\n\n    # convert into tokenized index\n    if start_position != -1 and end_position != -1:\n        start_position = original_to_tokenized_index[start_position]\n        end_position = original_to_tokenized_index[end_position]\n\n    # make sure at least one object in `examples`\n    examples = []\n    max_doc_len = max_seq_len - len(question_tokens) - 3  # [CLS], [SEP], [SEP]\n\n    # take chunks with a stride of `doc_stride`\n    for doc_start in range(0, len(all_doc_tokens), doc_stride):\n        doc_end = doc_start + max_doc_len\n        # if truncated document does not contain annotated range\n        if not (doc_start <= start_position and end_position <= doc_end):\n            start, end, label = -1, -1, 'unknown'\n        else:\n            start = start_position - doc_start + len(question_tokens) + 2\n            end = end_position - doc_start + len(question_tokens) + 2\n            label = class_label\n\n        assert -1 <= start < max_seq_len, f'start position is out of range: {start}'\n        assert -1 <= end < max_seq_len, f'end position is out of range: {end}'\n\n        doc_tokens = all_doc_tokens[doc_start:doc_end]\n        input_tokens = ['[CLS]'] + question_tokens + ['[SEP]'] + doc_tokens + ['[SEP]']\n        examples.append(\n            Example(\n                example_id=data['example_id'],\n                candidates=data['long_answer_candidates'],\n                annotations=annotations,\n                doc_start=doc_start,\n                question_len=len(question_tokens),\n                tokenized_to_original_index=tokenized_to_original_index,\n                input_ids=tokenizer.convert_tokens_to_ids(input_tokens),\n                start_position=start,\n                end_position=end,\n                class_label=label\n        ))\n\n    return examples","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.706651Z","iopub.status.idle":"2024-07-08T18:03:56.706976Z","shell.execute_reply.started":"2024-07-08T18:03:56.706804Z","shell.execute_reply":"2024-07-08T18:03:56.706817Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class JsonChunkReader(JsonReader):\n    \"\"\"JsonReader provides an interface for reading in a JSON file.\n    \"\"\"\n    \n    def __init__(\n        self,\n        filepath_or_buffer: str,\n        convert_data: Callable[[str], List[Example]],\n        orient: str = None,\n        typ: str = 'frame',\n        dtype: bool = None,\n        convert_axes: bool = None,\n        convert_dates: bool = True,\n        keep_default_dates: bool = True,\n        numpy: bool = False,\n        precise_float: bool = False,\n        date_unit: str = None,\n        encoding: str = None,\n        lines: bool = True,\n        chunksize: int = 2000,\n        compression: str = None,\n    ):\n        super(JsonChunkReader, self).__init__(\n            str(filepath_or_buffer),\n            orient=orient, typ=typ, dtype=dtype,\n            convert_axes=convert_axes,\n            convert_dates=convert_dates,\n            keep_default_dates=keep_default_dates,\n            numpy=numpy, precise_float=precise_float,\n            date_unit=date_unit, encoding=encoding,\n            lines=lines, chunksize=chunksize,\n            compression=compression\n        )\n        self.convert_data = convert_data\n        \n    def __next__(self):\n        lines = list(itertools.islice(self.data, self.chunksize))\n        if lines:\n            with Pool(2) as p:\n                obj = p.map(self.convert_data, lines)\n            return obj\n\n        self.close()\n        raise StopIteration","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.707839Z","iopub.status.idle":"2024-07-08T18:03:56.70818Z","shell.execute_reply.started":"2024-07-08T18:03:56.708022Z","shell.execute_reply":"2024-07-08T18:03:56.708037Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TextDataset(Dataset):\n    \"\"\"Dataset for [TensorFlow 2.0 Question Answering](https://www.kaggle.com/c/tensorflow2-question-answering).\n    \n    Parameters\n    ----------\n    examples : list of Example\n        The whole Dataset.\n    \"\"\"\n    \n    def __init__(self, examples: List[Example]):\n        self.examples = examples\n        \n    def __len__(self) -> int:\n        return len(self.examples)\n    \n    def __getitem__(self, index):\n        annotated = list(\n            filter(lambda example: example.class_label != 'unknown', self.examples[index]))\n        if len(annotated) == 0:\n            return random.choice(self.examples[index])\n        return random.choice(annotated)\n\n    \ndef collate_fn(examples: List[Example]) -> List[List[torch.Tensor]]:\n    # input tokens\n    max_len = max([len(example.input_ids) for example in examples])\n    tokens = np.zeros((len(examples), max_len), dtype=np.int64)\n    token_type_ids = np.ones((len(examples), max_len), dtype=np.int64)\n    for i, example in enumerate(examples):\n        row = example.input_ids\n        tokens[i, :len(row)] = row\n        token_type_id = [0 if i <= row.index(102) else 1\n                         for i in range(len(row))]  # 102 corresponds to [SEP]\n        token_type_ids[i, :len(row)] = token_type_id\n    attention_mask = tokens > 0\n    inputs = [torch.from_numpy(tokens),\n              torch.from_numpy(attention_mask),\n              torch.from_numpy(token_type_ids)]\n\n    # output labels\n    all_labels = ['long', 'no', 'short', 'unknown', 'yes']\n    start_positions = np.array([example.start_position for example in examples])\n    end_positions = np.array([example.end_position for example in examples])\n    class_labels = [all_labels.index(example.class_label) for example in examples]\n    start_positions = np.where(start_positions >= max_len, -1, start_positions)\n    end_positions = np.where(end_positions >= max_len, -1, end_positions)\n    labels = [torch.LongTensor(start_positions),\n              torch.LongTensor(end_positions),\n              torch.LongTensor(class_labels)]\n\n    return [inputs, labels]","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.709473Z","iopub.status.idle":"2024-07-08T18:03:56.709812Z","shell.execute_reply.started":"2024-07-08T18:03:56.709649Z","shell.execute_reply":"2024-07-08T18:03:56.709664Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BertForQuestionAnswering(BertPreTrainedModel):\n    \"\"\"BERT model for QA and classification tasks.\n    \n    Parameters\n    ----------\n    config : transformers.BertConfig. Configuration class for BERT.\n        \n    Returns\n    -------\n    start_logits : torch.Tensor with shape (batch_size, sequence_size).\n        Starting scores of each tokens.\n    end_logits : torch.Tensor with shape (batch_size, sequence_size).\n        Ending scores of each tokens.\n    classifier_logits : torch.Tensor with shape (batch_size, num_classes).\n        Classification scores of each labels.\n    \"\"\"\n\n    def __init__(self, config):\n        super(BertForQuestionAnswering, self).__init__(config)\n        self.bert = BertModel(config)\n        self.qa_outputs = nn.Linear(config.hidden_size, 2)  # start/end\n        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n        self.classifier = nn.Linear(config.hidden_size, config.num_labels)\n        self.init_weights()\n\n    def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):\n        outputs = self.bert(input_ids,\n                            attention_mask=attention_mask,\n                            token_type_ids=token_type_ids,\n                            position_ids=position_ids, \n                            head_mask=head_mask)\n\n        sequence_output = outputs[0]\n        pooled_output = outputs[1]\n\n        # predict start & end position\n        qa_logits = self.qa_outputs(sequence_output)\n        start_logits, end_logits = qa_logits.split(1, dim=-1)\n        start_logits = start_logits.squeeze(-1)\n        end_logits = end_logits.squeeze(-1)\n    \n        # classification\n        pooled_output = self.dropout(pooled_output)\n        classifier_logits = self.classifier(pooled_output)\n\n        return start_logits, end_logits, classifier_logits","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.711216Z","iopub.status.idle":"2024-07-08T18:03:56.711532Z","shell.execute_reply.started":"2024-07-08T18:03:56.711372Z","shell.execute_reply":"2024-07-08T18:03:56.711385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def loss_fn(preds, labels):\n    start_preds, end_preds, class_preds = preds\n    start_labels, end_labels, class_labels = labels\n    \n    start_loss = nn.CrossEntropyLoss(ignore_index=-1)(start_preds, start_labels)\n    end_loss = nn.CrossEntropyLoss(ignore_index=-1)(end_preds, end_labels)\n    class_loss = nn.CrossEntropyLoss()(class_preds, class_labels)\n    return start_loss + end_loss + class_loss","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.713128Z","iopub.status.idle":"2024-07-08T18:03:56.713439Z","shell.execute_reply.started":"2024-07-08T18:03:56.713285Z","shell.execute_reply":"2024-07-08T18:03:56.713298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = BertForQuestionAnswering.from_pretrained(bert_model, num_labels=5)\nmodel = model.to(device)\n\nparam_optimizer = list(model.named_parameters())\nno_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\noptimizer_grouped_parameters = [\n    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},\n    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]\nnum_train_optimization_steps = int(n_epochs * train_size / batch_size / accumulation_steps)\nnum_warmup_steps = int(num_train_optimization_steps * warmup)\n\noptimizer = AdamW(optimizer_grouped_parameters, lr=lr, correct_bias=False)\nscheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_train_optimization_steps)\n\nmodel, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)\nmodel.zero_grad()\nmodel = model.train()\n\ntokenizer = BertTokenizer.from_pretrained(bert_model, do_lower_case=do_lower_case)\nconvert_func = functools.partial(convert_data,\n                                 tokenizer=tokenizer,\n                                 max_seq_len=max_seq_len,\n                                 max_question_len=max_question_len,\n                                 doc_stride=doc_stride)\ndata_reader = JsonChunkReader(DATA_PATH, convert_func, chunksize=chunksize)\n\nglobal_step = 0\nfor examples in tqdm(data_reader, total=int(np.ceil(train_size / chunksize))):\n    train_dataset = TextDataset(examples)\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)\n    for x_batch, y_batch in train_loader:\n        x_batch, attention_mask, token_type_ids = x_batch\n        y_batch = (y.to(device) for y in y_batch)\n\n        y_pred = model(x_batch.to(device),\n                       attention_mask=attention_mask.to(device),\n                       token_type_ids=token_type_ids.to(device))\n        loss = loss_fn(y_pred, y_batch)\n        with amp.scale_loss(loss, optimizer) as scaled_loss:\n            scaled_loss.backward()\n        if (global_step + 1) % accumulation_steps == 0:\n            optimizer.step()\n            scheduler.step()\n            model.zero_grad()\n\n        global_step += 1\n        \n    if (time.time() - start_time) / 3600 > 7:\n        break\n\ndel examples, train_dataset, train_loader\ngc.collect()\n\ntorch.save(model.state_dict(), output_model_file)\ntorch.save(optimizer.state_dict(), output_optimizer_file)\ntorch.save(amp.state_dict(), output_amp_file)","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.714523Z","iopub.status.idle":"2024-07-08T18:03:56.714821Z","shell.execute_reply.started":"2024-07-08T18:03:56.714671Z","shell.execute_reply":"2024-07-08T18:03:56.714683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Check cuda version so we can install the correct version\n!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.716269Z","iopub.status.idle":"2024-07-08T18:03:56.716604Z","shell.execute_reply.started":"2024-07-08T18:03:56.716434Z","shell.execute_reply":"2024-07-08T18:03:56.716448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nprint(\"Pytorch version：\")\nprint(torch.__version__)\nprint(\"CUDA Version: \")\nprint(torch.version.cuda)\nprint(\"cuDNN version is :\")\nprint(torch.backends.cudnn.version())","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.717584Z","iopub.status.idle":"2024-07-08T18:03:56.717913Z","shell.execute_reply.started":"2024-07-08T18:03:56.717732Z","shell.execute_reply":"2024-07-08T18:03:56.717745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!wget https://download.pytorch.org/whl/cu111/torch-1.9.0%2Bcu111-cp37-cp37m-linux_x86_64.whl\n!wget https://download.pytorch.org/whl/cu111/torchvision-0.10.0%2Bcu111-cp37-cp37m-linux_x86_64.whl  ","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.719279Z","iopub.status.idle":"2024-07-08T18:03:56.719727Z","shell.execute_reply.started":"2024-07-08T18:03:56.719493Z","shell.execute_reply":"2024-07-08T18:03:56.719511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f'trained {global_step * batch_size} samples')\nprint(f'training time: {(time.time() - start_time) / 3600:.1f} hours')","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.721563Z","iopub.status.idle":"2024-07-08T18:03:56.722026Z","shell.execute_reply.started":"2024-07-08T18:03:56.721774Z","shell.execute_reply":"2024-07-08T18:03:56.721794Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def eval_collate_fn(examples: List[Example]) -> Tuple[List[torch.Tensor], List[Example]]:\n    # input tokens\n    max_len = max([len(example.input_ids) for example in examples])\n    tokens = np.zeros((len(examples), max_len), dtype=np.int64)\n    token_type_ids = np.ones((len(examples), max_len), dtype=np.int64)\n    for i, example in enumerate(examples):\n        row = example.input_ids\n        tokens[i, :len(row)] = row\n        token_type_id = [0 if i <= row.index(102) else 1\n                         for i in range(len(row))]  # 102 corresponds to [SEP]\n        token_type_ids[i, :len(row)] = token_type_id\n    attention_mask = tokens > 0\n    inputs = [torch.from_numpy(tokens),\n              torch.from_numpy(attention_mask),\n              torch.from_numpy(token_type_ids)]\n\n    return inputs, examples\n\n\ndef eval_model(\n    model: nn.Module,\n    valid_loader: DataLoader,\n    device: torch.device = torch.device('cuda')\n) -> Dict[str, float]:\n    \"\"\"Compute validation score.\n    \n    Parameters\n    ----------\n    model : nn.Module\n        Model for prediction.\n    valid_loader : DataLoader\n        Data loader of validation data.\n    device : torch.device, optional\n        Device for computation.\n    \n    Returns\n    -------\n    dict\n        Scores of validation data.\n        `long_score`: score of long answers\n        `short_score`: score of short answers\n        `overall_score`: score of the competition metric\n    \"\"\"\n    model.to(device)\n    model.eval()\n    with torch.no_grad():\n        result = Result()\n        for inputs, examples in tqdm(valid_loader):\n            input_ids, attention_mask, token_type_ids = inputs\n            y_preds = model(input_ids.to(device),\n                            attention_mask.to(device),\n                            token_type_ids.to(device))\n            \n            start_preds, end_preds, class_preds = (p.detach().cpu() for p in y_preds)\n            start_logits, start_index = torch.max(start_preds, dim=1)\n            end_logits, end_index = torch.max(end_preds, dim=1)\n\n            # span logits minus the cls logits seems to be close to the best\n            cls_logits = start_preds[:, 0] + end_preds[:, 0]  # '[CLS]' logits\n            logits = start_logits + end_logits - cls_logits  # (batch_size,)\n            indices = torch.stack((start_index, end_index)).transpose(0, 1)  # (batch_size, 2)\n            result.update(examples, logits.numpy(), indices.numpy(), class_preds.numpy())\n\n    return result.score()\n\n\nclass Result(object):\n    \"\"\"Stores results of all test data.\n    \"\"\"\n    \n    def __init__(self):\n        self.examples = {}\n        self.results = {}\n        self.best_scores = defaultdict(float)\n        self.class_labels = ['LONG', 'NO', 'SHORT', 'UNKNOWN', 'YES']\n        \n    @staticmethod\n    def is_valid_index(example: Example, index: List[int]) -> bool:\n        \"\"\"Return whether valid index or not.\n        \"\"\"\n        start_index, end_index = index\n        if start_index > end_index:\n            return False\n        if start_index <= example.question_len + 2:\n            return False\n        return True\n        \n    def update(\n        self,\n        examples: List[Example],\n        logits: torch.Tensor,\n        indices: torch.Tensor,\n        class_preds: torch.Tensor\n    ):\n        \"\"\"Update batch objects.\n        \n        Parameters\n        ----------\n        examples : list of Example\n        logits : np.ndarray with shape (batch_size,)\n            Scores of each examples..\n        indices : np.ndarray with shape (batch_size, 2)\n            `start_index` and `end_index` pairs of each examples.\n        class_preds : np.ndarray with shape (batch_size, num_classes)\n            Class predicition scores of each examples.\n        \"\"\"\n        for i, example in enumerate(examples):\n            if self.is_valid_index(example, indices[i]) and \\\n               self.best_scores[example.example_id] < logits[i]:\n                self.best_scores[example.example_id] = logits[i]\n                self.examples[example.example_id] = example\n                self.results[example.example_id] = [\n                    example.doc_start, indices[i], class_preds[i]]\n\n    def _generate_predictions(self) -> Generator[Dict, None, None]:\n        \"\"\"Generate predictions of each examples.\n        \"\"\"\n        for example_id in self.results.keys():\n            doc_start, index, class_pred = self.results[example_id]\n            example = self.examples[example_id]\n            tokenized_to_original_index = example.tokenized_to_original_index\n            short_start_index = tokenized_to_original_index[doc_start + index[0]]\n            short_end_index = tokenized_to_original_index[doc_start + index[1]]\n            long_start_index = -1\n            long_end_index = -1\n            for candidate in example.candidates:\n                if candidate['start_token'] <= short_start_index and \\\n                   short_end_index <= candidate['end_token']:\n                    long_start_index = candidate['start_token']\n                    long_end_index = candidate['end_token']\n                    break\n            yield {\n                'example': example,\n                'long_answer': [long_start_index, long_end_index],\n                'short_answer': [short_start_index, short_end_index],\n                'yes_no_answer': class_pred\n            }\n\n    def end(self) -> Dict[str, Dict]:\n        \"\"\"Get predictions in submission format.\n        \"\"\"\n        preds = {}\n        for pred in self._generate_predictions():\n            example = pred['example']\n            long_start_index, long_end_index = pred['long_answer']\n            short_start_index, short_end_index = pred['short_answer']\n            class_pred = pred['yes_no_answer']\n\n            long_answer = f'{long_start_index}:{long_end_index}' if long_start_index != -1 else np.nan\n            short_answer = f'{short_start_index}:{short_end_index}'\n            class_pred = self.class_labels[class_pred.argmax()]\n            short_answer += ' ' + class_pred if class_pred in ['YES', 'NO'] else ''\n            preds[f'{example.example_id}_long'] = long_answer\n            preds[f'{example.example_id}_short'] = short_answer\n        return preds\n\n    def score(self) -> Dict[str, float]:\n        \"\"\"Calculate score of all examples.\n        \"\"\"\n\n        def _safe_divide(x: int, y: int) -> float:\n            \"\"\"Compute x / y, but return 0 if y is zero.\n            \"\"\"\n            if y == 0:\n                return 0.\n            else:\n                return x / y\n\n        def _compute_f1(answer_stats: List[List[bool]]) -> float:\n            \"\"\"Computes F1, precision, recall for a list of answer scores.\n            \"\"\"\n            has_answer, has_pred, is_correct = list(zip(*answer_stats))\n            precision = _safe_divide(sum(is_correct), sum(has_pred))\n            recall = _safe_divide(sum(is_correct), sum(has_answer))\n            f1 = _safe_divide(2 * precision * recall, precision + recall)\n            return f1\n\n        long_scores = []\n        short_scores = []\n        for pred in self._generate_predictions():\n            example = pred['example']\n            long_pred = pred['long_answer']\n            short_pred = pred['short_answer']\n            class_pred = pred['yes_no_answer']\n            yes_no_label = self.class_labels[class_pred.argmax()]\n\n            # long score\n            long_label = example.annotations['long_answer']\n            has_answer = long_label['candidate_index'] != -1\n            has_pred = long_pred[0] != -1 and long_pred[1] != -1\n            is_correct = False\n            if long_label['start_token'] == long_pred[0] and \\\n               long_label['end_token'] == long_pred[1]:\n                is_correct = True\n            long_scores.append([has_answer, has_pred, is_correct])\n\n            # short score\n            short_labels = example.annotations['short_answers']\n            class_pred = example.annotations['yes_no_answer']\n            has_answer = yes_no_label != 'NONE' or len(short_labels) != 0\n            has_pred = class_pred != 'NONE' or (short_pred[0] != -1 and short_pred[1] != -1)\n            is_correct = False\n            if class_pred in ['YES', 'NO']:\n                is_correct = yes_no_label == class_pred\n            else:\n                for short_label in short_labels:\n                    if short_label['start_token'] == short_pred[0] and \\\n                       short_label['end_token'] == short_pred[1]:\n                        is_correct = True\n                        break\n            short_scores.append([has_answer, has_pred, is_correct])\n\n        long_score = _compute_f1(long_scores)\n        short_score = _compute_f1(short_scores)\n        return {\n            'long_score': long_score,\n            'short_score': short_score,\n            'overall_score': (long_score + short_score) / 2\n        }","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.723536Z","iopub.status.idle":"2024-07-08T18:03:56.723999Z","shell.execute_reply.started":"2024-07-08T18:03:56.723753Z","shell.execute_reply":"2024-07-08T18:03:56.723773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"eval_start_time = time.time()\n\nvalid_data = next(data_reader)\nvalid_data = list(itertools.chain.from_iterable(valid_data))\nvalid_dataset = Subset(valid_data, range(len(valid_data)))\nvalid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, collate_fn=eval_collate_fn)\nvalid_scores = eval_model(model, valid_loader, device=device)\n\nprint(f'calculate validation score done in {(time.time() - eval_start_time) / 60:.1f} minutes.')","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.725346Z","iopub.status.idle":"2024-07-08T18:03:56.725794Z","shell.execute_reply.started":"2024-07-08T18:03:56.72556Z","shell.execute_reply":"2024-07-08T18:03:56.72558Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"long_score = valid_scores['long_score']\nshort_score = valid_scores['short_score']\noverall_score = valid_scores['overall_score']\nprint('validation scores:')\nprint(f'\\tlong score    : {long_score:.4f}')\nprint(f'\\tshort score   : {short_score:.4f}')\nprint(f'\\toverall score : {overall_score:.4f}')\nprint(f'all process done in {(time.time() - start_time) / 3600:.1f} hours.')","metadata":{"execution":{"iopub.status.busy":"2024-07-08T18:03:56.7272Z","iopub.status.idle":"2024-07-08T18:03:56.727519Z","shell.execute_reply.started":"2024-07-08T18:03:56.727359Z","shell.execute_reply":"2024-07-08T18:03:56.727372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}