import os
from collections import defaultdict
import argparse
import json
import string
import random
import numpy as np
import torch
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, Dataset
from transformers import RobertaTokenizer, AutoModel, AutoConfig, AutoModelWithLMHead
from scripts.triviaqa_utils import evaluation_utils
import pytorch_lightning as pl
from pytorch_lightning.logging import TestTubeLogger
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel
from longformer.longformer import Longformer
from longformer.sliding_chunks import pad_to_window_size
class TriviaQADataset(Dataset):
"""
Largely based on
https://github.com/allenai/allennlp/blob/master/allennlp/data/dataset_readers/reading_comprehension/triviaqa.py
and
https://github.com/huggingface/transformers/blob/master/examples/run_squad.py
"""
def __init__(self, file_path, tokenizer, max_seq_len, max_doc_len, doc_stride,
max_num_answers, ignore_seq_with_no_answers, max_question_len):
assert os.path.isfile(file_path)
self.file_path = file_path
with open(self.file_path, "r", encoding='utf-8') as f:
print(f'reading file: {self.file_path}')
self.data_json = json.load(f)['data']
print(f'done reading file: {self.file_path}')
self.tokenizer = tokenizer
self.max_seq_len = max_seq_len
self.max_doc_len = max_doc_len
self.doc_stride = doc_stride
self.max_num_answers = max_num_answers
self.ignore_seq_with_no_answers = ignore_seq_with_no_answers
self.max_question_len = max_question_len
# A mapping from qid to an int, which can be synched across gpus using `torch.distributed`
if 'train' not in self.file_path: # only for the evaluation set
self.val_qid_string_to_int_map = \
{
self._get_qid(entry["paragraphs"][0]['qas'][0]['id']): index
for index, entry in enumerate(self.data_json)
}
else:
self.val_qid_string_to_int_map = None
def _normalize_text(self, text: str) -> str: # copied from the official triviaqa repo
return " ".join(
[
token
for token in text.lower().strip(self.STRIPPED_CHARACTERS).split()
if token not in self.IGNORED_TOKENS
]
)
IGNORED_TOKENS = {"a", "an", "the"}
STRIPPED_CHARACTERS = string.punctuation + "".join([u"‘", u"’", u"´", u"`", "_"])
def __len__(self):
return len(self.data_json)
def __getitem__(self, idx):
entry = self.data_json[idx]
tensors_list = self.one_example_to_tensors(entry, idx)
assert len(tensors_list) == 1
return tensors_list[0]
def one_example_to_tensors(self, example, idx):
def is_whitespace(c):
if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
return True
return False
tensors_list = []
for paragraph in example["paragraphs"]:
paragraph_text = paragraph["context"]
doc_tokens = []
char_to_word_offset = []
prev_is_whitespace = True
for c in paragraph_text:
if is_whitespace(c):
prev_is_whitespace = True
else:
if prev_is_whitespace:
doc_tokens.append(c)
else:
doc_tokens[-1] += c
prev_is_whitespace = False
char_to_word_offset.append(len(doc_tokens) - 1)
for qa in paragraph["qas"]:
question_text = qa["question"]
start_position = None
end_position = None
orig_answer_text = None
answer_spans = []
for answer in qa["answers"]:
orig_answer_text = answer["text"]
answer_offset = answer["answer_start"]
answer_length = len(orig_answer_text)
try:
start_position = char_to_word_offset[answer_offset]
end_position = char_to_word_offset[answer_offset + answer_length - 1]
token_ids = self.tokenizer.encode(orig_answer_text)
except RuntimeError:
print(f'Reading example {idx} failed')
start_position = 0
end_position = 0
answer_spans.append({'start': start_position, 'end': end_position,
'text': orig_answer_text, 'token_ids': token_ids})
# ===== Given an example, convert it into tensors =============
query_tokens = self.tokenizer.tokenize(question_text)
query_tokens = query_tokens[:self.max_question_len]
tok_to_orig_index = []
orig_to_tok_index = []
all_doc_tokens = []
for (i, token) in enumerate(doc_tokens):
orig_to_tok_index.append(len(all_doc_tokens))
# hack: the line below should have been `self.tokenizer.tokenize(token')`
# but roberta tokenizer uses a different subword if the token is the beginning of the string
# or in the middle. So for all tokens other than the first, simulate that it is not the first
# token by prepending a period before tokenizing, then dropping the period afterwards
sub_tokens = self.tokenizer.tokenize(f'. {token}')[1:] if i > 0 else self.tokenizer.tokenize(token)
for sub_token in sub_tokens:
tok_to_orig_index.append(i)
all_doc_tokens.append(sub_token)
all_doc_tokens = all_doc_tokens[:self.max_doc_len]
# The -3 accounts for [CLS], [SEP] and [SEP]
max_tokens_per_doc_slice = self.max_seq_len - len(query_tokens) - 3
assert max_tokens_per_doc_slice > 0
if self.doc_stride < 0:
# negative doc_stride indicates no sliding window, but using first slice
self.doc_stride = -100 * len(all_doc_tokens) # large -ve value for the next loop to execute once
input_ids_list = []
input_mask_list = []
segment_ids_list = []
start_positions_list = []
end_positions_list = []
answer_token_ids_list = []
for slice_start in range(0, len(all_doc_tokens), max_tokens_per_doc_slice - self.doc_stride):
slice_end = min(slice_start + max_tokens_per_doc_slice, len(all_doc_tokens))
doc_slice_tokens = all_doc_tokens[slice_start:slice_end]
tokens = [self.tokenizer.cls_token] + query_tokens + [self.tokenizer.sep_token] \
+ doc_slice_tokens + [self.tokenizer.sep_token]
segment_ids = [0] * (len(query_tokens) + 2) + [1] * (len(doc_slice_tokens) + 1)
assert len(segment_ids) == len(tokens)
input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
input_mask = [1] * len(input_ids)
if self.doc_stride >= 0: # no need to pad if document is not strided
# Zero-pad up to the sequence length.
padding_len = self.max_seq_len - len(input_ids)
input_ids.extend([self.tokenizer.pad_token_id] * padding_len)
input_mask.extend([0] * padding_len)
segment_ids.extend([0] * padding_len)
assert len(input_ids) == self.max_seq_len
assert len(input_mask) == self.max_seq_len
assert len(segment_ids) == self.max_seq_len
doc_offset = len(query_tokens) + 2 - slice_start
start_positions = []
end_positions = []
answer_token_ids = []
for answer_span in answer_spans:
start_position = answer_span['start']
end_position = answer_span['end']
tok_start_position_in_doc = orig_to_tok_index[start_position]
not_end_of_doc = int(end_position + 1 < len(orig_to_tok_index))
tok_end_position_in_doc = orig_to_tok_index[end_position + not_end_of_doc] - not_end_of_doc
if tok_start_position_in_doc < slice_start or tok_end_position_in_doc > slice_end:
# this answer is outside the current slice
continue
start_positions.append(tok_start_position_in_doc + doc_offset)
end_positions.append(tok_end_position_in_doc + doc_offset)
answer_token_ids.append(answer_span['token_ids'])
assert len(start_positions) == len(end_positions)
if self.ignore_seq_with_no_answers and len(start_positions) == 0:
continue
# answers from start_positions and end_positions if > self.max_num_answers
start_positions = start_positions[:self.max_num_answers]
end_positions = end_positions[:self.max_num_answers]
answer_token_ids = answer_token_ids[:self.max_num_answers]
# -1 padding up to self.max_num_answers
padding_len = self.max_num_answers - len(start_positions)
start_positions.extend([-1] * padding_len)
end_positions.extend([-1] * padding_len)
answer_token_ids.extend([[]] * padding_len)
# replace duplicate start/end positions with `-1` because duplicates can result into -ve loss values
found_start_positions = set()
found_end_positions = set()
found_answer_token_ids = set()
for i, (start_position, end_position, answer_tokens) in enumerate(
zip(start_positions, end_positions, answer_token_ids)
):
if start_position in found_start_positions:
start_positions[i] = -1
if end_position in found_end_positions:
end_positions[i] = -1
answer_tokens_as_str = ','.join([str(x) for x in answer_tokens])
if answer_tokens_as_str in found_answer_token_ids:
answer_token_ids[i] = []
found_start_positions.add(start_position)
found_end_positions.add(end_position)
found_answer_token_ids.add(answer_tokens_as_str)
input_ids_list.append(input_ids)
input_mask_list.append(input_mask)
segment_ids_list.append(segment_ids)
start_positions_list.append(start_positions)
end_positions_list.append(end_positions)
answer_token_ids_list.append(answer_token_ids)
# pad answers in answer_token_ids_list to the longest answer
max_answer_len = max([len(item) for sublist in answer_token_ids_list for item in sublist]) # flat list
if max_answer_len == 0:
max_answer_len = 2
for answers_of_one_slice in answer_token_ids_list:
for answer_tokens in answers_of_one_slice:
if len(answer_tokens) == 0:
# TODO: