{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip install ../input/apex01cp37cp37mlinux-x86-64whl/apex-0.1-cp37-cp37m-linux_x86_64.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport sys\nimport collections\nimport json\n\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# pip list","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nfrom tqdm import tqdm\nimport torch.nn as nn\nfrom torch import optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torch\nimport re\nimport json\nfrom transformers import BertTokenizer, BertModel, BertPreTrainedModel, BertConfig, AlbertTokenizer, AlbertModel, AlbertPreTrainedModel, AlbertConfig\nimport time\nfrom apex import amp\n\npath='/kaggle/input/models/models'\nos.chdir(path)\nos.listdir(path)\n\ndef reduce1(n_candidate=10, th_candidate=0.2):\n\n    class TFQADataset(Dataset):\n        def __init__(self, id_list):\n            self.id_list=id_list \n        def __len__(self):\n            return len(self.id_list)\n        def __getitem__(self, index):\n            return self.id_list[index]\n\n    class Collator(object):\n        def __init__(self, data_dict, new_token_dict, tokenizer, max_seq_len=384, max_question_len=64):\n            self.data_dict = data_dict\n            self.new_token_dict = new_token_dict\n            self.tokenizer = tokenizer\n            self.max_seq_len = max_seq_len\n            self.max_question_len = max_question_len\n\n        def _get_input_ids(self, doc_id, candidate_index):\n            data = self.data_dict[doc_id]\n            question_tokens = self.tokenizer.tokenize(data['question_text'])[:self.max_question_len]\n            doc_words = data['document_text'].split()\n\n            max_answer_tokens = self.max_seq_len-len(question_tokens)-3 # [CLS],[SEP],[SEP]\n            candidate = data['long_answer_candidates'][candidate_index]\n            candidate_start = candidate['start_token']\n            candidate_end = candidate['end_token']\n            candidate_words = doc_words[candidate_start:candidate_end]  \n            # Loop through to add html tokens as new tokens here\n            for i, word in enumerate(candidate_words):\n                if re.match(r'<.+>', word):\n                    if word in self.new_token_dict: \n                        candidate_words[i] = self.new_token_dict[word]\n                    else:\n                        candidate_words[i] = '<'     \n\n            words_to_tokens_index = []\n            candidate_tokens = []\n            for i, word in enumerate(candidate_words):\n                words_to_tokens_index.append(len(candidate_tokens))\n                tokens = self.tokenizer.tokenize(word)\n                if len(candidate_tokens)+len(tokens) > max_answer_tokens: # token length cannot be longer than the global max length (360)\n                    break\n                for token in tokens:\n                    candidate_tokens.append(token)\n\n            input_tokens = ['[CLS]'] + question_tokens + ['[SEP]'] + candidate_tokens + ['[SEP]']\n            input_ids = self.tokenizer.convert_tokens_to_ids(input_tokens)\n\n            return input_ids, len(input_ids)\n        \n        def __call__(self, batch_ids):\n            batch_size = len(batch_ids)\n\n            # pre-compute all the input within the batch without padding to determine the actual batch max sequence length\n            batch_input_ids_temp = []\n            batch_seq_len = []\n\n            for i, (doc_id, candidate_index) in enumerate(batch_ids):\n                input_ids, seq_len = self._get_input_ids(doc_id, candidate_index)\n                batch_input_ids_temp.append(input_ids)\n                batch_seq_len.append(seq_len)\n\n            batch_max_seq_len = max(batch_seq_len) # set max sequence length to be the maximun length in a batch, to save computation \n            batch_input_ids = np.zeros((batch_size, batch_max_seq_len), dtype=np.int64)\n            batch_token_type_ids = np.ones((batch_size, batch_max_seq_len), dtype=np.int64)\n\n            for i in range(batch_size):\n                input_ids = batch_input_ids_temp[i]\n                batch_input_ids[i, :len(input_ids)] = input_ids\n                batch_token_type_ids[i, :len(input_ids)] = [0 if k<=input_ids.index(102) else 1 for k in range(len(input_ids))]\n\n            batch_attention_mask = batch_input_ids > 0\n\n            return torch.from_numpy(batch_input_ids),torch.from_numpy(batch_attention_mask),torch.from_numpy(batch_token_type_ids)\n\n    # https://www.kaggle.com/sakami/tfqa-pytorch-baseline\n    class BertForQuestionAnswering(BertPreTrainedModel):\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\n\n\n    # prepare input\n    json_dir = '../../../input/tensorflow2-question-answering/simplified-nq-test.jsonl'\n\n    # The id_candidate_list keeps all the combination of document ids and their candidates. So we essentially run predictions on all the candidates.\n    id_candidate_list = []\n    # Store the lengths of each candidate for rearranging based on candidate length, this can help improve inference speed significantly.\n    id_candidate_len_list = [] \n    # Keep a dictionary for length checking.\n    id_candidate_len_dict = {}\n    # list of document ids.\n    id_list = []\n    # Keeps the texts and candidates.\n    data_dict = {}\n    # for debugging only\n    max_data = 9999999999\n    with open(json_dir) as f:\n        for n, line in tqdm(enumerate(f)):\n            if n > max_data:\n                break\n            data = json.loads(line)\n            data_id = data['example_id']\n            id_list.append(data_id)\n\n            # initialize data_dict\n            data_dict[data_id] = {\n                                  'document_text': data['document_text'],\n                                  'question_text': data['question_text'], \n                                  'long_answer_candidates': data['long_answer_candidates'],                \n                                 }\n        \n            question_len = len(data['question_text'].split())\n\n            # We use the wite space tokenzied version to estimate candidate length here.\n            for i in range(len(data['long_answer_candidates'])):\n                id_candidate_list.append((data_id, i))\n                candidate_len = question_len+data['long_answer_candidates'][i]['end_token']-data['long_answer_candidates'][i]['start_token']\n                id_candidate_len_list.append(candidate_len)\n                id_candidate_len_dict[(data_id, i)] = candidate_len\n\n    print(len(id_candidate_list))\n\n    # Sort based on the length of each candidate.\n    id_candidate_len_list = np.array(id_candidate_len_list)\n    sorted_index = np.argsort(id_candidate_len_list)\n    id_candidate_list_sorted = []\n    for i in range(len(id_candidate_list)):\n        id_candidate_list_sorted.append(id_candidate_list[sorted_index[i]])\n\n\n    # hyperparameters\n    max_seq_len = 360\n    max_question_len = 64\n    batch_size = int(768/4)\n\n    # build model\n    model_path = '../models/models/bert/'\n    config = BertConfig.from_pretrained(model_path)\n    config.num_labels = 5\n    config.vocab_size = 30531\n    tokenizer = BertTokenizer.from_pretrained(model_path, do_lower_case=True)\n    model = BertForQuestionAnswering.from_pretrained('../models/models/bert/', config=config)\n\n    # add new tokens\n    new_token_dict = {\n                      '<P>':'qw1',\n                      '<Table>':'qw2',\n                      '<Tr>':'qw3',\n                      '<Ul>':'qw4',\n                      '<Ol>':'qw5',\n                      '<Fl>':'qw6',\n                      '<Li>':'qw7',\n                      '<Dd>':'qw8',\n                      '<Dt>':'qw9',\n                     }\n    new_token_list = [\n                      'qw1',\n                      'qw2',\n                      'qw3',\n                      'qw4',\n                      'qw5',\n                      'qw6',\n                      'qw7',\n                      'qw8',\n                      'qw9',\n                     ]\n\n    num_added_toks = tokenizer.add_tokens(new_token_list)\n    # print('We have added', num_added_toks, 'tokens')\n    model.resize_token_embeddings(len(tokenizer))\n\n\n    model.cuda()\n    optimizer = optim.Adam(model.parameters(), lr=1e-5)\n    model, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\",verbosity=0)\n    if torch.cuda.device_count() > 1:\n        model = torch.nn.DataParallel(model)\n\n\n    # testing\n\n    # iterator for testing\n    test_datagen = TFQADataset(id_list=id_candidate_list_sorted)\n    test_collate = Collator(data_dict=data_dict, \n                            new_token_dict=new_token_dict,\n                            tokenizer=tokenizer, \n                            max_seq_len=max_seq_len, \n                            max_question_len=max_question_len)\n    test_generator = DataLoader(dataset=test_datagen,\n                                collate_fn=test_collate,\n                                batch_size=batch_size,\n                                shuffle=False,\n                                num_workers=16,\n                                pin_memory=True)\n\n\n    model.eval()\n    test_prob3 = np.zeros((len(id_candidate_list_sorted),5),dtype=np.float32) # class\n    for j,(batch_input_ids, batch_attention_mask, batch_token_type_ids) in tqdm(enumerate(test_generator)):\n        with torch.no_grad():\n            start = j*batch_size\n            end = start+batch_size\n            if j == len(test_generator)-1:\n                end = len(test_generator.dataset)\n            batch_input_ids = batch_input_ids.cuda()\n            batch_attention_mask = batch_attention_mask.cuda()\n            batch_token_type_ids = batch_token_type_ids.cuda()\n            # We don't need the span output here.\n            _, _, logits3 = model(batch_input_ids, batch_attention_mask, batch_token_type_ids)\n            test_prob3[start:end] += F.softmax(logits3,dim=1).cpu().data.numpy()\n\n\n    # initialize a temp dictionary\n    temp_dict = {}\n    for doc_id in id_list:\n        temp_dict[doc_id] = np.zeros((len(data_dict[doc_id]['long_answer_candidates']),),dtype=np.float32)\n\n    # input long answer probs into the temp dictionary\n    for i, (doc_id, candidate_index) in tqdm(enumerate(id_candidate_list_sorted)):\n        temp_dict[doc_id][candidate_index] = 1.0 - test_prob3[i,0] # 1-no_answer_score\n\n    # get list of survived id-candidates\n    id_candidate_list1 = []\n    id_candidate_len_list1 = []\n    for doc_id in tqdm(id_list):\n        long_prob_array = temp_dict[doc_id].copy()\n        sorted_index = np.argsort(long_prob_array)[::-1]\n        count = 0\n        for n in range(len(sorted_index)):\n            if count>=n_candidate:\n                break\n            else:\n                if temp_dict[doc_id][sorted_index[n]]>th_candidate:\n                    id_candidate_list1.append((doc_id, sorted_index[n]))\n                    id_candidate_len_list1.append(id_candidate_len_dict[(doc_id, sorted_index[n])])\n                    count += 1\n\n    # sort and return\n    sorted_index = np.argsort(id_candidate_len_list1)\n    id_candidate_list_sorted1 = []\n    for i in range(len(id_candidate_list1)):\n        id_candidate_list_sorted1.append(id_candidate_list1[sorted_index[i]])\n\n    print(len(id_candidate_list_sorted1))\n\n    return data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted1\n\n\ndef reduce2(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, n_candidate=10, th_candidate=0.2):\n\n    class TFQADataset(Dataset):\n        def __init__(self, id_list):\n            self.id_list=id_list \n        def __len__(self):\n            return len(self.id_list)\n        def __getitem__(self, index):\n            return self.id_list[index]\n\n    class Collator(object):\n        def __init__(self, data_dict, new_token_dict, tokenizer, max_seq_len=384, max_question_len=64):\n            self.data_dict = data_dict\n            self.new_token_dict = new_token_dict\n            self.tokenizer = tokenizer\n            self.max_seq_len = max_seq_len\n            self.max_question_len = max_question_len\n\n        def _get_input_ids(self, doc_id, candidate_index):\n            data = self.data_dict[doc_id]\n            question_tokens = self.tokenizer.tokenize(data['question_text'])[:self.max_question_len]\n            doc_words = data['document_text'].split()\n\n            max_answer_tokens = self.max_seq_len-len(question_tokens)-3 # [CLS],[SEP],[SEP]\n            candidate = data['long_answer_candidates'][candidate_index]\n            candidate_start = candidate['start_token']\n            candidate_end = candidate['end_token']\n            candidate_words = doc_words[candidate_start:candidate_end]  \n            for i, word in enumerate(candidate_words):\n                if re.match(r'<.+>', word):\n                    if word in self.new_token_dict: \n                        candidate_words[i] = self.new_token_dict[word]\n                    else:\n                        candidate_words[i] = '<'     \n\n            words_to_tokens_index = []\n            candidate_tokens = []\n            for i, word in enumerate(candidate_words):\n                words_to_tokens_index.append(len(candidate_tokens))\n                tokens = self.tokenizer.tokenize(word)\n                if len(candidate_tokens)+len(tokens) > max_answer_tokens:\n                    break\n                for token in tokens:\n                    candidate_tokens.append(token)\n\n            input_tokens = ['[CLS]'] + question_tokens + ['[SEP]'] + candidate_tokens + ['[SEP]']\n            input_ids = self.tokenizer.convert_tokens_to_ids(input_tokens)\n\n            return input_ids, len(input_ids)\n        \n        def __call__(self, batch_ids):\n            batch_size = len(batch_ids)\n\n            batch_input_ids_temp = []\n            batch_seq_len = []\n\n            for i, (doc_id, candidate_index) in enumerate(batch_ids):\n                input_ids, seq_len = self._get_input_ids(doc_id, candidate_index)\n                batch_input_ids_temp.append(input_ids)\n                batch_seq_len.append(seq_len)\n\n            batch_max_seq_len = max(batch_seq_len)\n            batch_input_ids = np.zeros((batch_size, batch_max_seq_len), dtype=np.int64)\n            batch_token_type_ids = np.ones((batch_size, batch_max_seq_len), dtype=np.int64)\n\n            for i in range(batch_size):\n                input_ids = batch_input_ids_temp[i]\n                batch_input_ids[i, :len(input_ids)] = input_ids\n                batch_token_type_ids[i, :len(input_ids)] = [0 if k<=input_ids.index(102) else 1 for k in range(len(input_ids))]\n\n            batch_attention_mask = batch_input_ids > 0\n\n            return torch.from_numpy(batch_input_ids),torch.from_numpy(batch_attention_mask),torch.from_numpy(batch_token_type_ids)\n\n\n    class BertForQuestionAnswering(BertPreTrainedModel):\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\n\n\n    # hyperparameters\n    max_seq_len = 360\n    max_question_len = 64\n    batch_size = int(384/4)\n\n\n    # build model\n    model_path = '../models/models/bert/'\n    config = BertConfig.from_pretrained(model_path)\n    config.num_labels = 5\n    config.vocab_size = 30531\n    tokenizer = BertTokenizer.from_pretrained(model_path, do_lower_case=True)\n    model = BertForQuestionAnswering.from_pretrained('../models/models/bert/', config=config)\n\n    # add new tokens\n    new_token_dict = {\n                      '<P>':'qw1',\n                      '<Table>':'qw2',\n                      '<Tr>':'qw3',\n                      '<Ul>':'qw4',\n                      '<Ol>':'qw5',\n                      '<Fl>':'qw6',\n                      '<Li>':'qw7',\n                      '<Dd>':'qw8',\n                      '<Dt>':'qw9',\n                     }\n    new_token_list = [\n                      'qw1',\n                      'qw2',\n                      'qw3',\n                      'qw4',\n                      'qw5',\n                      'qw6',\n                      'qw7',\n                      'qw8',\n                      'qw9',\n                     ]\n\n    num_added_toks = tokenizer.add_tokens(new_token_list)\n    print('We have added', num_added_toks, 'tokens')\n    model.resize_token_embeddings(len(tokenizer))\n\n\n    model.cuda()\n    optimizer = optim.Adam(model.parameters(), lr=1e-5)\n    model, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\",verbosity=0)\n    if torch.cuda.device_count() > 1:\n        model = torch.nn.DataParallel(model)\n\n\n    # testing\n\n    # iterator for testing\n    test_datagen = TFQADataset(id_list=id_candidate_list_sorted)\n    test_collate = Collator(data_dict=data_dict, \n                            new_token_dict=new_token_dict,\n                            tokenizer=tokenizer, \n                            max_seq_len=max_seq_len, \n                            max_question_len=max_question_len)\n    test_generator = DataLoader(dataset=test_datagen,\n                                collate_fn=test_collate,\n                                batch_size=batch_size,\n                                shuffle=False,\n                                num_workers=16,\n                                pin_memory=True)\n\n\n    model.eval()\n    test_prob3 = np.zeros((len(id_candidate_list_sorted),5),dtype=np.float32) # class\n    for j,(batch_input_ids, batch_attention_mask, batch_token_type_ids) in tqdm(enumerate(test_generator)):\n        with torch.no_grad():\n            start = j*batch_size\n            end = start+batch_size\n            if j == len(test_generator)-1:\n                end = len(test_generator.dataset)\n            batch_input_ids = batch_input_ids.cuda()\n            batch_attention_mask = batch_attention_mask.cuda()\n            batch_token_type_ids = batch_token_type_ids.cuda()\n            _, _, logits3 = model(batch_input_ids, batch_attention_mask, batch_token_type_ids)\n            test_prob3[start:end] += F.softmax(logits3,dim=1).cpu().data.numpy()\n\n\n    # initialize a temp dictionary\n    temp_dict = {}\n    for doc_id in id_list:\n        temp_dict[doc_id] = np.zeros((len(data_dict[doc_id]['long_answer_candidates']),),dtype=np.float32)\n\n    # input long answer probs into the temp dictionary\n    for i, (doc_id, candidate_index) in tqdm(enumerate(id_candidate_list_sorted)):\n        temp_dict[doc_id][candidate_index] = 1.0 - test_prob3[i,0] # 1-no_answer_score\n\n    # get list of survived id-candidates\n    id_candidate_list1 = []\n    id_candidate_len_list1 = []\n    for doc_id in tqdm(id_list):\n        long_prob_array = temp_dict[doc_id].copy()\n        sorted_index = np.argsort(long_prob_array)[::-1]\n        count = 0\n        for n in range(len(sorted_index)):\n            if count>=n_candidate:\n                break\n            else:\n                if temp_dict[doc_id][sorted_index[n]]>th_candidate:\n                    id_candidate_list1.append((doc_id, sorted_index[n]))\n                    id_candidate_len_list1.append(id_candidate_len_dict[(doc_id, sorted_index[n])])\n                    count += 1\n\n    # sort and return\n    sorted_index = np.argsort(id_candidate_len_list1)\n    id_candidate_list_sorted1 = []\n    for i in range(len(id_candidate_list1)):\n        id_candidate_list_sorted1.append(id_candidate_list1[sorted_index[i]])\n\n    print(len(id_candidate_list_sorted1))\n\n    return id_candidate_list_sorted1\n\n\ndef albert_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len):\n\n    class TFQADataset(Dataset):\n        def __init__(self, id_list):\n            self.id_list=id_list \n        def __len__(self):\n            return len(self.id_list)\n        def __getitem__(self, index):\n            return self.id_list[index]\n\n    class Collator(object):\n        def __init__(self, data_dict, new_token_dict, tokenizer, max_seq_len=384, max_question_len=64):\n            self.data_dict = data_dict\n            self.new_token_dict = new_token_dict\n            self.tokenizer = tokenizer\n            self.max_seq_len = max_seq_len\n            self.max_question_len = max_question_len\n\n        def _get_input_ids(self, doc_id, candidate_index):\n            data = self.data_dict[doc_id]\n            question_tokens = self.tokenizer.tokenize(data['question_text'])[:self.max_question_len]\n            doc_words = data['document_text'].split()\n\n            max_answer_tokens = self.max_seq_len-len(question_tokens)-3 # [CLS],[SEP],[SEP]\n            candidate = data['long_answer_candidates'][candidate_index]\n            candidate_start = candidate['start_token']\n            candidate_end = candidate['end_token']\n            candidate_words = doc_words[candidate_start:candidate_end]  \n            for i, word in enumerate(candidate_words):\n                if re.match(r'<.+>', word):\n                    if word in self.new_token_dict: \n                        candidate_words[i] = self.new_token_dict[word]\n                    else:\n                        candidate_words[i] = 'qw99'    \n\n            words_to_tokens_index = []\n            tokens_to_words_index = []\n            candidate_tokens = []\n            for i, word in enumerate(candidate_words):\n                words_to_tokens_index.append(len(candidate_tokens))\n                tokens = self.tokenizer.tokenize(word)\n                if len(candidate_tokens)+len(tokens) > max_answer_tokens:\n                    break\n                for token in tokens:\n                    tokens_to_words_index.append(i)\n                    candidate_tokens.append(token)\n\n            input_tokens = ['[CLS]'] + question_tokens + ['[SEP]'] + candidate_tokens + ['[SEP]']\n            input_ids = self.tokenizer.convert_tokens_to_ids(input_tokens)\n\n            return input_ids, words_to_tokens_index, len(input_ids), len(question_tokens)+2\n        \n        def __call__(self, batch_ids):\n            batch_size = len(batch_ids)\n\n            batch_input_ids_temp = []\n            batch_seq_len = []\n\n            batch_offset = []\n            batch_words_to_tokens_index = []\n\n            for i, (doc_id, candidate_index) in enumerate(batch_ids):\n                input_ids, words_to_tokens_index, seq_len, offset = self._get_input_ids(doc_id, candidate_index)\n                batch_input_ids_temp.append(input_ids)\n                batch_seq_len.append(seq_len)\n                batch_offset.append(offset)\n                batch_words_to_tokens_index.append(words_to_tokens_index)\n\n            batch_max_seq_len = max(batch_seq_len)\n            batch_input_ids = np.zeros((batch_size, batch_max_seq_len), dtype=np.int64)\n            batch_token_type_ids = np.ones((batch_size, batch_max_seq_len), dtype=np.int64)\n\n            for i in range(batch_size):\n                input_ids = batch_input_ids_temp[i]\n                batch_input_ids[i, :len(input_ids)] = input_ids\n                batch_token_type_ids[i, :len(input_ids)] = [0 if k<=input_ids.index(3) else 1 for k in range(len(input_ids))]\n\n            batch_attention_mask = batch_input_ids > 0\n\n            return torch.from_numpy(batch_input_ids), torch.from_numpy(batch_attention_mask), torch.from_numpy(batch_token_type_ids), batch_words_to_tokens_index, batch_offset, batch_max_seq_len\n\n\n    class AlbertForQuestionAnswering(AlbertPreTrainedModel):\n\n        def __init__(self, config):\n            super(AlbertForQuestionAnswering, self).__init__(config)\n            self.albert = AlbertModel(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.albert(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\n\n\n    # hyperparameters\n    max_seq_len = 360\n    max_question_len = 64\n    batch_size = int(128/4)\n\n\n    # build model\n    model_path = '../models/models/albert/'\n    config = AlbertConfig.from_pretrained(model_path)\n    config.num_labels = 5\n    config.vocab_size = 30010\n    tokenizer = AlbertTokenizer.from_pretrained(model_path, do_lower_case=True)\n    model = AlbertForQuestionAnswering.from_pretrained(model_dir, config=config)\n\n    # add new tokens\n    new_token_dict = {\n                      '<P>':'qw1',\n                      '<Table>':'qw2',\n                      '<Tr>':'qw3',\n                      '<Ul>':'qw4',\n                      '<Ol>':'qw5',\n                      '<Fl>':'qw6',\n                      '<Li>':'qw7',\n                      '<Dd>':'qw8',\n                      '<Dt>':'qw9',\n                     }\n    new_token_list = [\n                      'qw1',\n                      'qw2',\n                      'qw3',\n                      'qw4',\n                      'qw5',\n                      'qw6',\n                      'qw7',\n                      'qw8',\n                      'qw9',\n                      'qw99',\n                     ]\n\n    num_added_toks = tokenizer.add_tokens(new_token_list)\n    # print('We have added', num_added_toks, 'tokens')\n    model.resize_token_embeddings(len(tokenizer))\n\n\n    model.cuda()\n    optimizer = optim.Adam(model.parameters(), lr=1e-5)\n    model, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\",verbosity=0)\n    if torch.cuda.device_count() > 1:\n        model = torch.nn.DataParallel(model)\n\n\n    # testing\n\n    # iterator for testing\n    test_datagen = TFQADataset(id_list=id_candidate_list_sorted)\n    test_collate = Collator(data_dict=data_dict, \n                            new_token_dict=new_token_dict,\n                            tokenizer=tokenizer, \n                            max_seq_len=max_seq_len, \n                            max_question_len=max_question_len)\n    test_generator = DataLoader(dataset=test_datagen,\n                                collate_fn=test_collate,\n                                batch_size=batch_size,\n                                shuffle=False,\n                                num_workers=16,\n                                pin_memory=True)\n\n\n    model.eval()\n    list_offset = []\n    list_words_to_tokens_index = []\n    test_prob1 = np.zeros((len(id_candidate_list_sorted),max_seq_len),dtype=np.float32) # start\n    test_prob2 = np.zeros((len(id_candidate_list_sorted),max_seq_len),dtype=np.float32) # end\n    test_prob3 = np.zeros((len(id_candidate_list_sorted),5),dtype=np.float32) # class\n    for j,(batch_input_ids, batch_attention_mask, batch_token_type_ids, batch_words_to_tokens_index, batch_offset, batch_max_seq_len) in tqdm(enumerate(test_generator)):\n        with torch.no_grad():\n            start = j*batch_size\n            end = start+batch_size\n            if j == len(test_generator)-1:\n                end = len(test_generator.dataset)\n            batch_input_ids = batch_input_ids.cuda()\n            batch_attention_mask = batch_attention_mask.cuda()\n            batch_token_type_ids = batch_token_type_ids.cuda()\n            logits1, logits2, logits3 = model(batch_input_ids, batch_attention_mask, batch_token_type_ids)\n            test_prob1[start:end, :batch_max_seq_len] += F.softmax(logits1,dim=1).cpu().data.numpy()\n            test_prob2[start:end, :batch_max_seq_len] += F.softmax(logits2,dim=1).cpu().data.numpy()\n            test_prob3[start:end] += F.softmax(logits3,dim=1).cpu().data.numpy()\n            list_words_to_tokens_index += batch_words_to_tokens_index\n            list_offset += batch_offset\n\n    test_word_prob1 = np.zeros((len(id_candidate_list_sorted),word_len),dtype=np.float32) # start\n    test_word_prob2 = np.zeros((len(id_candidate_list_sorted),word_len),dtype=np.float32) # end\n    for i in range(len(id_candidate_list_sorted)):\n        for j in range(len(list_words_to_tokens_index[i])):\n            test_word_prob1[i,j] = test_prob1[i, list_words_to_tokens_index[i][j]+list_offset[i]]\n            test_word_prob2[i,j] = test_prob2[i, list_words_to_tokens_index[i][j]+list_offset[i]]\n\n\n    return test_word_prob1, test_word_prob2, test_prob3\n\n\n# This function performs a full prediction on the validation set using a fast model (bert-base) to reduce the candidates for larger model predictions.\n# Propose only the top-k (top10 in this case) most probable candidates from each document, each candidate must have long answer probability larger than a threshold (0.2), the rest candidates are set to negative.\n# id_candidate_list_sorted stores (document id, candidate number) as keys, each of each contains its long answer probability (score).\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def bert_large_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len):\n\n    class TFQADataset(Dataset):\n        def __init__(self, id_list):\n            self.id_list=id_list \n        def __len__(self):\n            return len(self.id_list)\n        def __getitem__(self, index):\n            return self.id_list[index]\n\n    class Collator(object):\n        def __init__(self, data_dict, new_token_dict, tokenizer, max_seq_len=384, max_question_len=64):\n            self.data_dict = data_dict\n            self.new_token_dict = new_token_dict\n            self.tokenizer = tokenizer\n            self.max_seq_len = max_seq_len\n            self.max_question_len = max_question_len\n\n        def _get_input_ids(self, doc_id, candidate_index):\n            data = self.data_dict[doc_id]\n            question_tokens = self.tokenizer.tokenize(data['question_text'])[:self.max_question_len]\n            doc_words = data['document_text'].split()\n\n            max_answer_tokens = self.max_seq_len-len(question_tokens)-3 # [CLS],[SEP],[SEP]\n            candidate = data['long_answer_candidates'][candidate_index]\n            candidate_start = candidate['start_token']\n            candidate_end = candidate['end_token']\n            candidate_words = doc_words[candidate_start:candidate_end]  \n            for i, word in enumerate(candidate_words):\n                if re.match(r'<.+>', word):\n                    if word in self.new_token_dict: \n                        candidate_words[i] = self.new_token_dict[word]\n                    else:\n                        candidate_words[i] = '<'     \n\n            words_to_tokens_index = []\n            candidate_tokens = []\n            for i, word in enumerate(candidate_words):\n                words_to_tokens_index.append(len(candidate_tokens))\n                tokens = self.tokenizer.tokenize(word)\n                if len(candidate_tokens)+len(tokens) > max_answer_tokens:\n                    break\n                for token in tokens:\n                    candidate_tokens.append(token)\n\n            input_tokens = ['[CLS]'] + question_tokens + ['[SEP]'] + candidate_tokens + ['[SEP]']\n            input_ids = self.tokenizer.convert_tokens_to_ids(input_tokens)\n\n            return input_ids, words_to_tokens_index, len(input_ids), len(question_tokens)+2\n        \n        def __call__(self, batch_ids):\n            batch_size = len(batch_ids)\n\n            batch_input_ids_temp = []\n            batch_seq_len = []\n\n            batch_offset = []\n            batch_words_to_tokens_index = []\n\n            for i, (doc_id, candidate_index) in enumerate(batch_ids):\n                input_ids, words_to_tokens_index, seq_len, offset = self._get_input_ids(doc_id, candidate_index)\n                batch_input_ids_temp.append(input_ids)\n                batch_seq_len.append(seq_len)\n                batch_offset.append(offset)\n                batch_words_to_tokens_index.append(words_to_tokens_index)\n\n            batch_max_seq_len = max(batch_seq_len)\n            batch_input_ids = np.zeros((batch_size, batch_max_seq_len), dtype=np.int64)\n            batch_token_type_ids = np.ones((batch_size, batch_max_seq_len), dtype=np.int64)\n\n            for i in range(batch_size):\n                input_ids = batch_input_ids_temp[i]\n                batch_input_ids[i, :len(input_ids)] = input_ids\n                batch_token_type_ids[i, :len(input_ids)] = [0 if k<=input_ids.index(102) else 1 for k in range(len(input_ids))]\n\n            batch_attention_mask = batch_input_ids > 0\n\n            return torch.from_numpy(batch_input_ids), torch.from_numpy(batch_attention_mask), torch.from_numpy(batch_token_type_ids), batch_words_to_tokens_index, batch_offset, batch_max_seq_len\n\n\n    class BertForQuestionAnswering(BertPreTrainedModel):\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\n\n\n    # hyperparameters\n    max_seq_len = 360\n    max_question_len = 64\n    batch_size = int(384/4)\n\n\n    # build model\n    model_path = '../models/models/bert/'\n    config = BertConfig.from_pretrained(model_path)\n    config.num_labels = 5\n    config.vocab_size = 30531\n    tokenizer = BertTokenizer.from_pretrained(model_path, do_lower_case=True)\n    model = BertForQuestionAnswering.from_pretrained(model_dir, config=config)\n\n    # add new tokens\n    new_token_dict = {\n                      '<P>':'qw1',\n                      '<Table>':'qw2',\n                      '<Tr>':'qw3',\n                      '<Ul>':'qw4',\n                      '<Ol>':'qw5',\n                      '<Fl>':'qw6',\n                      '<Li>':'qw7',\n                      '<Dd>':'qw8',\n                      '<Dt>':'qw9',\n                     }\n    new_token_list = [\n                      'qw1',\n                      'qw2',\n                      'qw3',\n                      'qw4',\n                      'qw5',\n                      'qw6',\n                      'qw7',\n                      'qw8',\n                      'qw9',\n                     ]\n\n    num_added_toks = tokenizer.add_tokens(new_token_list)\n    # print('We have added', num_added_toks, 'tokens')\n    model.resize_token_embeddings(len(tokenizer))\n\n\n    model.cuda()\n    optimizer = optim.Adam(model.parameters(), lr=1e-5)\n    model, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\",verbosity=0)\n    if torch.cuda.device_count() > 1:\n        model = torch.nn.DataParallel(model)\n\n\n    # testing\n\n    # iterator for testing\n    test_datagen = TFQADataset(id_list=id_candidate_list_sorted)\n    test_collate = Collator(data_dict=data_dict, \n                            new_token_dict=new_token_dict,\n                            tokenizer=tokenizer, \n                            max_seq_len=max_seq_len, \n                            max_question_len=max_question_len)\n    test_generator = DataLoader(dataset=test_datagen,\n                                collate_fn=test_collate,\n                                batch_size=batch_size,\n                                shuffle=False,\n                                num_workers=16,\n                                pin_memory=True)\n\n\n    model.eval()\n    list_offset = []\n    list_words_to_tokens_index = []\n    test_prob1 = np.zeros((len(id_candidate_list_sorted),max_seq_len),dtype=np.float32) # start\n    test_prob2 = np.zeros((len(id_candidate_list_sorted),max_seq_len),dtype=np.float32) # end\n    test_prob3 = np.zeros((len(id_candidate_list_sorted),5),dtype=np.float32) # class\n    for j,(batch_input_ids, batch_attention_mask, batch_token_type_ids, batch_words_to_tokens_index, batch_offset, batch_max_seq_len) in tqdm(enumerate(test_generator)):\n        with torch.no_grad():\n            start = j*batch_size\n            end = start+batch_size\n            if j == len(test_generator)-1:\n                end = len(test_generator.dataset)\n            batch_input_ids = batch_input_ids.cuda()\n            batch_attention_mask = batch_attention_mask.cuda()\n            batch_token_type_ids = batch_token_type_ids.cuda()\n            logits1, logits2, logits3 = model(batch_input_ids, batch_attention_mask, batch_token_type_ids)\n            test_prob1[start:end, :batch_max_seq_len] += F.softmax(logits1,dim=1).cpu().data.numpy()\n            test_prob2[start:end, :batch_max_seq_len] += F.softmax(logits2,dim=1).cpu().data.numpy()\n            test_prob3[start:end] += F.softmax(logits3,dim=1).cpu().data.numpy()\n            list_words_to_tokens_index += batch_words_to_tokens_index\n            list_offset += batch_offset\n\n    # From token-level to word-level span predictions. Use the first token of each word for word-level representation.\n    test_word_prob1 = np.zeros((len(id_candidate_list_sorted),word_len),dtype=np.float32) # start\n    test_word_prob2 = np.zeros((len(id_candidate_list_sorted),word_len),dtype=np.float32) # end\n    for i in range(len(id_candidate_list_sorted)):\n        for j in range(len(list_words_to_tokens_index[i])):\n            test_word_prob1[i,j] = test_prob1[i, list_words_to_tokens_index[i][j]+list_offset[i]]\n            test_word_prob2[i,j] = test_prob2[i, list_words_to_tokens_index[i][j]+list_offset[i]]\n\n    return test_word_prob1, test_word_prob2, test_prob3\n\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"start_time = time.time()\ndata_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted = reduce1(n_candidate=10, th_candidate=0.2)\nprint(\"--- %s seconds ---\" % (time.time() - start_time))\n\n# Futher reduce the number of candidates from top10 to top4.\nstart_time = time.time()\nid_candidate_list_sorted = reduce2(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, n_candidate=4, th_candidate=0.35)\nprint(\"--- %s seconds ---\" % (time.time() - start_time))\n\n# Acutual predictions start here. \nstart_time = time.time()\n# We keep word-level posterior start and end vectors. Since the token-level length is set to 360, this number should be enough for word-level.\nword_len = 360\n# Initialize start and end \"\"word-level\"\" prob vectors for easier probability averaging between different models equiped with different tokenizers.\nstart_prob = np.zeros((len(id_candidate_list_sorted),word_len),dtype=np.float32)\nend_prob = np.zeros((len(id_candidate_list_sorted),word_len),dtype=np.float32)\nstart_label = np.zeros((len(id_candidate_list_sorted),),dtype=int)\nend_label = np.zeros((len(id_candidate_list_sorted),),dtype=int)\n# class_prob stores the 5-class classifier prob outputs.\n# no answer(0), long but not short answer(1), short answer with span(2), NO(3), YES(4)\nclass_prob = np.zeros((len(id_candidate_list_sorted),5),dtype=np.float32)\n\n# Perform prediction using two albert-xxl and two bert-large models. Weighted average of both long and short predictions for ensembling.\n# model_dir = '../albert-xxlarge-v2_2/weights/epoch2/'\n# test_prob1, test_prob2, test_prob3 = albert_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len)\n# start_prob += 0.3*test_prob1\n# end_prob += 0.3*test_prob2\n# class_prob += 0.3*test_prob3\n# model_dir = '../albert-xxlarge-v2_3/weights/epoch2/'\n# test_prob1, test_prob2, test_prob3 = albert_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len)\n# start_prob += 0.3*test_prob1\n# end_prob += 0.3*test_prob2\n# class_prob += 0.3*test_prob3\n# model_dir = '../bert-large-uncased_4/weights/epoch3/'\n# test_prob1, test_prob2, test_prob3 = bert_large_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len)\n# start_prob += 0.2*test_prob1\n# end_prob += 0.2*test_prob2\n# class_prob += 0.2*test_prob3\n\nprint(\"before bert predicting\")\n\nweight_albert = 0.8\n\nmodel_dir = '../models/models/bert/'\ntest_prob1, test_prob2, test_prob3 = bert_large_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len)\nstart_prob += (1-weight_albert)*test_prob1\nend_prob += (1-weight_albert)*test_prob2\nclass_prob += (1-weight_albert)*test_prob3\n\nprint(\"before albert predicting\")\n\nmodel_dir = '../models/models/albert/'\ntest_prob1, test_prob2, test_prob3 = albert_predict(data_dict, id_list, id_candidate_len_dict, id_candidate_list_sorted, model_dir, word_len)\nstart_prob += weight_albert*test_prob1\nend_prob += weight_albert*test_prob2\nclass_prob += weight_albert*test_prob3\n\n# print(\"before saving\")\n\n# The start and end words have the largest probabilities.\nstart_label = np.argmax(start_prob, axis=1)\nend_label = np.argmax(end_prob, axis=1)\n\n# initialize a temporary dictionary to store prediction values.\ntemp_dict = {}\nfor doc_id in id_list:\n    temp_dict[doc_id] = {\n                         'long_answer': {'start_token': -1, 'end_token': -1},\n                         'long_answer_score': -1.0,\n                         'short_answers': [{'start_token': -1, 'end_token': -1}],\n                         'short_answers_score': -1.0,\n                         'yes_no_answer': 'NONE'\n                        }\n\n# from cadidates to document\nfor i, (doc_id, candidate_index) in tqdm(enumerate(id_candidate_list_sorted)):\n    # process long answer\n    long_answer_score = 1.0 - class_prob[i,0] # 1 - no_answer_score\n    if long_answer_score > temp_dict[doc_id]['long_answer_score']:\n        temp_dict[doc_id]['long_answer_score'] = long_answer_score\n        temp_dict[doc_id]['long_answer']['start_token'] = data_dict[doc_id]['long_answer_candidates'][candidate_index]['start_token']\n        temp_dict[doc_id]['long_answer']['end_token'] = data_dict[doc_id]['long_answer_candidates'][candidate_index]['end_token']\n        # process short answer\n        short_answer_score = 1.0 - class_prob[i,0] - class_prob[i,1] # 1 - no_answer_score - long_but_not_short_answer_score\n        temp_dict[doc_id]['short_answers_score'] = short_answer_score\n\n        temp_dict[doc_id]['short_answers'][0]['start_token'] = -1\n        temp_dict[doc_id]['short_answers'][0]['end_token'] = -1\n        temp_dict[doc_id]['yes_no_answer'] = 'NONE'\n        if max([class_prob[i,3], class_prob[i,4]]) > class_prob[i,2]:\n            if class_prob[i,3] > class_prob[i,4]:\n                temp_dict[doc_id]['yes_no_answer'] = 'NO'\n            else:\n                temp_dict[doc_id]['yes_no_answer'] = 'YES'\n        else:\n            short_start_word = int(start_label[i]) + data_dict[doc_id]['long_answer_candidates'][candidate_index]['start_token']\n            short_end_word = int(end_label[i]) + data_dict[doc_id]['long_answer_candidates'][candidate_index]['start_token']\n            if short_end_word > short_start_word:\n                temp_dict[doc_id]['short_answers'][0]['start_token'] = short_start_word\n                temp_dict[doc_id]['short_answers'][0]['end_token'] = short_end_word\n\n# Copy the temporary dictionary into the final dictionary that meets the required format for validation.\nfinal_dict = {}\nfinal_dict['predictions'] = []\nfor doc_id in id_list:\n    prediction_dict = {\n                       'example_id': doc_id,\n                       'long_answer': {'start_byte': -1, 'end_byte': -1, 'start_token': temp_dict[doc_id]['long_answer']['start_token'], 'end_token': temp_dict[doc_id]['long_answer']['end_token']},\n                       'long_answer_score': temp_dict[doc_id]['long_answer_score'],\n                       'short_answers': [{'start_byte': -1, 'end_byte': -1, 'start_token': temp_dict[doc_id]['short_answers'][0]['start_token'], 'end_token': temp_dict[doc_id]['short_answers'][0]['end_token']}],\n                       'short_answers_score': temp_dict[doc_id]['short_answers_score'],\n                       'yes_no_answer': temp_dict[doc_id]['yes_no_answer']\n                      }\n    final_dict['predictions'].append(prediction_dict)\n\n# with open('predictions.json', 'w') as fp:\n#     json.dump(final_dict, fp)\nprint(\"--- %s seconds ---\" % (time.time() - start_time))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"answers_df = pd.DataFrame.from_dict(final_dict)\nlong_best_threshold = 0.4228079319000244-0.1\nshort_best_threshold = 0.5448758341372013-0.1\n# long_best_threshold = 0.5211206674575806-0.1\n# short_best_threshold = 0.5261547937989235-0.1\n\ndef df_long_index_score(df):\n    answers = []\n    cont = 0\n    for i in range(len(df)):\n        index = {}\n        if df[\"predictions\"][i]['long_answer_score'] > long_best_threshold: \n            index['start'] = df[\"predictions\"][i]['long_answer']['start_token']\n            index['end'] = df[\"predictions\"][i]['long_answer']['end_token']\n            index['score'] = df[\"predictions\"][i]['long_answer_score']\n            index = [index]\n            answers.append(index)\n        else:\n            answers.append([])\n    return answers\n\n\ndef df_short_index_score(df):\n    answers = []\n    cont = 0\n    for i in range(len(df)):\n        index = {}\n        if df[\"predictions\"][i]['short_answers_score'] > short_best_threshold and df[\"predictions\"][i]['short_answers'][0]['start_token'] != -1:\n            index['start'] = df[\"predictions\"][i]['short_answers'][0]['start_token']\n            index['end'] = df[\"predictions\"][i]['short_answers'][0]['end_token']\n            index['score'] = df[\"predictions\"][i]['short_answers_score']\n            index = [index]\n            answers.append(index)\n        else:\n            answers.append([])\n    return answers\n\n\ndef df_example_id(df):\n    return df['example_id']\n\ndef create_answer(entry):\n    answer = []\n    for e in entry:\n        answer.append(str(e['start']) + ':' + str(e['end']))\n    if not answer:\n        answer = \"\"\n    return \", \".join(answer)\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# answers_df['answer'] = answers_df['predictions'].apply(df_long_index_score)\nanswers_df['long_indexes_and_scores'] = df_long_index_score(answers_df)\nanswers_df['short_indexes_and_scores'] = df_short_index_score(answers_df)\nanswers_df['example_id'] = answers_df['predictions'].apply(df_example_id)\n\n\n# list(answers_df[answers_df['example_id'].isin([\"-332839753184669166\"])][\"predictions\"])[0]['example_id']\n# answers_df[\"predictions\"][0]['short_answers'][0]['start_token']\n\n# answers_df[answers_df['example_id'].isin([\"-332839753184669166\"])] # [\"predictions\"]\n\nanswers_df = answers_df.drop(['predictions'], axis=1)\n\n# answers_df.head()\n\nanswers_df[\"long_answer\"] = answers_df['long_indexes_and_scores'].apply(\n    create_answer)\nanswers_df[\"short_answer\"] = answers_df['short_indexes_and_scores'].apply(\n    create_answer)\nanswers_df[\"example_id\"] = answers_df['example_id'].apply(lambda q: str(q))\n\nlong_answers = dict(zip(answers_df[\"example_id\"], answers_df[\"long_answer\"]))\nshort_answers = dict(zip(answers_df[\"example_id\"], answers_df[\"short_answer\"]))\n\n# answers_df.head()\n\nanswers_df = answers_df.drop(\n    ['long_indexes_and_scores', 'short_indexes_and_scores'], axis=1)\n# answers_df.head()\n\nsample_submission = pd.read_csv(\n    \"../../../input/tensorflow2-question-answering/sample_submission.csv\")\n\nlong_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\n    \"_long\")].apply(lambda q: long_answers[q[\"example_id\"].replace(\"_long\", \"\")], axis=1)\nshort_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\n    \"_short\")].apply(lambda q: short_answers[q[\"example_id\"].replace(\"_short\", \"\")], axis=1)\n\nsample_submission.loc[sample_submission[\"example_id\"].str.contains(\n    \"_long\"), \"PredictionString\"] = long_prediction_strings\nsample_submission.loc[sample_submission[\"example_id\"].str.contains(\n    \"_short\"), \"PredictionString\"] = short_prediction_strings\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"path = '/kaggle/working'\nos.chdir(path)\n\nsample_submission.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfor dirname, _, filenames in os.walk('/kaggle/working'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","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":4}