2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license store the tokenizers which insert 2 separators tokens returns tokenized answer spans that better match the annotated answer tokanswertext jointokenizer tokenizeoriganswertext for newstart in rangeinputstart inputend 1 for newend in rangeinputend newstart 1 1 textspan joindoctokensnewstart newend 1 if textspan tokanswertext return newstart newend return inputstart inputend def checkismaxcontextdocspans curspanindex position check if this is the max context doc span for the token if lendocspans 1 return true bestscore none bestspanindex none for spanindex docspan in enumeratedocspans end docspanstart docspanlength 1 if position docspanstart continue if position end continue numleftcontext position docspanstart numrightcontext end position score minnumleftcontext numrightcontext 0 01 docspanlength if bestscore is none or score bestscore bestscore score bestspanindex spanindex return curspanindex bestspanindex def iswhitespacec if c or c t or c r or c n or ordc 0x202f return true return false def squadconvertexampletofeatures example maxseqlength docstride maxquerylength paddingstrategy istraining features if istraining and not example isimpossible get start and end position startposition example startposition endposition example endposition if the answer cannot be found in the text then skip this example actualtext joinexample doctokensstartposition endposition 1 cleanedanswertext joinwhitespacetokenizeexample answertext if actualtext findcleanedanswertext 1 logger warningfcould not find answer actualtext vs cleanedanswertext return toktoorigindex origtotokindex alldoctokens for i token in enumerateexample doctokens origtotokindex appendlenalldoctokens if tokenizer class name in robertatokenizer longformertokenizer barttokenizer robertatokenizerfast longformertokenizerfast barttokenizerfast subtokens tokenizer tokenizetoken addprefixspacetrue else subtokens tokenizer tokenizetoken for subtoken in subtokens toktoorigindex appendi alldoctokens appendsubtoken if istraining and not example isimpossible tokstartposition origtotokindexexample startposition if example endposition lenexample doctokens 1 tokendposition origtotokindexexample endposition 1 1 else tokendposition lenalldoctokens 1 tokstartposition tokendposition improveanswerspan alldoctokens tokstartposition tokendposition tokenizer example answertext spans truncatedquery tokenizer encode example questiontext addspecialtokensfalse truncationtrue maxlengthmaxquerylength tokenizers who insert 2 sep tokens inbetween context question need to have special handling in the way they compute mask of added tokens tokenizertype typetokenizer name replacetokenizer lower sequenceaddedtokens tokenizer modelmaxlength tokenizer maxlensinglesentence 1 if tokenizertype in multiseptokenstokenizersset else tokenizer modelmaxlength tokenizer maxlensinglesentence sequencepairaddedtokens tokenizer modelmaxlength tokenizer maxlensentencespair spandoctokens alldoctokens while lenspans docstride lenalldoctokens define the side we want to truncate pad and the textpair sorting if tokenizer paddingside right texts truncatedquery pairs spandoctokens truncation truncationstrategy onlysecond value else texts spandoctokens pairs truncatedquery truncation truncationstrategy onlyfirst value encodeddict tokenizer encodeplus todothom update this logic texts pairs truncationtruncation paddingpaddingstrategy maxlengthmaxseqlength returnoverflowingtokenstrue stridemaxseqlength docstride lentruncatedquery sequencepairaddedtokens returntokentypeidstrue paragraphlen min lenalldoctokens lenspans docstride maxseqlength lentruncatedquery sequencepairaddedtokens if tokenizer padtokenid in encodeddictinputids if tokenizer paddingside right nonpaddedids encodeddictinputids encodeddictinputids indextokenizer padtokenid else lastpaddingidposition lenencodeddictinputids 1 encodeddictinputids 1 indextokenizer padtokenid nonpaddedids encodeddictinputidslastpaddingidposition 1 else nonpaddedids encodeddictinputids tokens tokenizer convertidstotokensnonpaddedids tokentoorigmap for i in rangeparagraphlen index lentruncatedquery sequenceaddedtokens i if tokenizer paddingside right else i tokentoorigmapindex toktoorigindexlenspans docstride i encodeddictparagraphlen paragraphlen encodeddicttokens tokens encodeddicttokentoorigmap tokentoorigmap encodeddicttruncatedquerywithspecialtokenslength lentruncatedquery sequenceaddedtokens encodeddicttokenismaxcontext encodeddictstart lenspans docstride encodeddictlength paragraphlen spans appendencodeddict if overflowingtokens not in encodeddict or overflowingtokens in encodeddict and lenencodeddictoverflowingtokens 0 break spandoctokens encodeddictoverflowingtokens for docspanindex in rangelenspans for j in rangespansdocspanindexparagraphlen ismaxcontext newcheckismaxcontextspans docspanindex docspanindex docstride j index j if tokenizer paddingside left else spansdocspanindextruncatedquerywithspecialtokenslength j spansdocspanindextokenismaxcontextindex ismaxcontext for span in spans identify the position of the cls token clsindex spaninputids indextokenizer clstokenid pmask mask with 1 for token than cannot be in the answer 0 for token which can be in an answer original tf implementation also keep the classification token set to 0 pmask np oneslikespantokentypeids if tokenizer paddingside right pmasklentruncatedquery sequenceaddedtokens 0 else pmasklenspantokens lentruncatedquery sequenceaddedtokens 0 padtokenindices np wherespaninputids tokenizer padtokenid specialtokenindices np asarray tokenizer getspecialtokensmaskspaninputids alreadyhasspecialtokenstrue nonzero pmaskpadtokenindices 1 pmaskspecialtokenindices 1 set the cls index to 0 the cls index can be used for impossible answers pmaskclsindex 0 spanisimpossible example isimpossible startposition 0 endposition 0 if istraining and not spanisimpossible for training if our document chunk does not contain an annotation we throw it out since there is nothing to predict docstart spanstart docend spanstart spanlength 1 outofspan false if not tokstartposition docstart and tokendposition docend outofspan true if outofspan startposition clsindex endposition clsindex spanisimpossible true else if tokenizer paddingside left docoffset 0 else docoffset lentruncatedquery sequenceaddedtokens startposition tokstartposition docstart docoffset endposition tokendposition docstart docoffset features append squadfeatures spaninputids spanattentionmask spantokentypeids clsindex pmask tolist exampleindex0 can not set uniqueid and exampleindex here they will be set after multiple processing uniqueid0 paragraphlenspanparagraphlen tokenismaxcontextspantokenismaxcontext tokensspantokens tokentoorigmapspantokentoorigmap startpositionstartposition endpositionendposition isimpossiblespanisimpossible qasidexample qasid return features def squadconvertexampletofeaturesinittokenizerforconvert pretrainedtokenizerbase global tokenizer tokenizer tokenizerforconvert def squadconvertexamplestofeatures examples tokenizer maxseqlength docstride maxquerylength istraining paddingstrategymaxlength returndatasetfalse threads1 tqdmenabledtrue defining helper methods convert to tensors and build dataset why have we split the batch into a tuple pytorch just has a list of tensors processor for the squad data set overridden by squadv1processor and squadv2processor used by the version 1 1 and version 2 0 of squad respectively creates a list of data processors squad squadexample using a tfds dataset args dataset the tfds dataset loaded from tensorflowdatasets loadsquad evaluate boolean specifying if in evaluation mode or in training mode returns list of squadexample examples python import tensorflowdatasets as tfds dataset tfds loadsquad trainingexamples getexamplesfromdatasetdataset evaluatefalse evaluationexamples getexamplesfromdatasetdataset evaluatetrue if evaluate dataset datasetvalidation else dataset datasettrain examples for tensordict in tqdmdataset examples appendself getexamplefromtensordicttensordict evaluateevaluate return examples def gettrainexamplesself datadir filenamenone if datadir is none datadir if self trainfile is none raise valueerrorsquadprocessor should be instantiated via squadv1processor or squadv2processor with open os path joindatadir self trainfile if filename is none else filename r encodingutf8 as reader inputdata json loadreaderdata return self createexamplesinputdata train def getdevexamplesself datadir filenamenone if datadir is none datadir if self devfile is none raise valueerrorsquadprocessor should be instantiated via squadv1processor or squadv2processor with open os path joindatadir self devfile if filename is none else filename r encodingutf8 as reader inputdata json loadreaderdata return self createexamplesinputdata dev def createexamplesself inputdata settype istraining settype train examples for entry in tqdminputdata title entrytitle for paragraph in entryparagraphs contexttext paragraphcontext for qa in paragraphqas qasid qaid questiontext qaquestion startpositioncharacter none answertext none answers isimpossible qa getisimpossible false if not isimpossible if istraining answer qaanswers0 answertext answertext startpositioncharacter answeranswerstart else answers qaanswers example squadexample qasidqasid questiontextquestiontext contexttextcontexttext answertextanswertext startpositioncharacterstartpositioncharacter titletitle isimpossibleisimpossible answersanswers examples appendexample return examples class squadv1processorsquadprocessor trainfile trainv1 1 json devfile devv1 1 json class squadv2processorsquadprocessor trainfile trainv2 0 json devfile devv2 0 json class squadexample def init self qasid questiontext contexttext answertext startpositioncharacter title answers isimpossiblefalse self qasid qasid self questiontext questiontext self contexttext contexttext self answertext answertext self title title self isimpossible isimpossible self answers answers self startposition self endposition 0 0 doctokens chartowordoffset previswhitespace true split on whitespace so that different tokens may be attributed to their original position for c in self contexttext if iswhitespacec previswhitespace true else if previswhitespace doctokens appendc else doctokens1 c previswhitespace false chartowordoffset appendlendoctokens 1 self doctokens doctokens self chartowordoffset chartowordoffset start and end positions only has a value during evaluation if startpositioncharacter is not none and not isimpossible self startposition chartowordoffsetstartpositioncharacter self endposition chartowordoffset minstartpositioncharacter lenanswertext 1 lenchartowordoffset 1 class squadfeatures def init self inputids attentionmask tokentypeids clsindex pmask exampleindex uniqueid paragraphlen tokenismaxcontext tokens tokentoorigmap startposition endposition isimpossible qasid str none encoding batchencoding none self inputids inputids self attentionmask attentionmask self tokentypeids tokentypeids self clsindex clsindex self pmask pmask self exampleindex exampleindex self uniqueid uniqueid self paragraphlen paragraphlen self tokenismaxcontext tokenismaxcontext self tokens tokens self tokentoorigmap tokentoorigmap self startposition startposition self endposition endposition self isimpossible isimpossible self qasid qasid self encoding encoding class squadresult def initself uniqueid startlogits endlogits starttopindexnone endtopindexnone clslogitsnone self startlogits startlogits self endlogits endlogits self uniqueid uniqueid if starttopindex self starttopindex starttopindex self endtopindex endtopindex self clslogits clslogits 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license store the tokenizers which insert 2 separators tokens returns tokenized answer spans that better match the annotated answer check if this is the max context doc span for the token check if this is the max context doc span for the token if len doc_spans 1 return true get start and end position if the answer cannot be found in the text then skip this example tokenizers who insert 2 sep tokens in between context question need to have special handling in the way they compute mask of added tokens define the side we want to truncate pad and the text pair sorting todo thom update this logic identify the position of the cls token p_mask mask with 1 for token than cannot be in the answer 0 for token which can be in an answer original tf implementation also keep the classification token set to 0 set the cls index to 0 the cls index can be used for impossible answers for training if our document chunk does not contain an annotation we throw it out since there is nothing to predict can not set unique_id and example_index here they will be set after multiple processing converts a list of examples into a list of features that can be directly given as input to a model it is model dependant and takes advantage of many of the tokenizer s features to create the model s inputs args examples list of data processors squad squadexample tokenizer an instance of a child of pretrainedtokenizer max_seq_length the maximum sequence length of the inputs doc_stride the stride used when the context is too large and is split across several features max_query_length the maximum length of the query is_training whether to create features for model evaluation or model training padding_strategy default to max_length which padding strategy to use return_dataset default false either pt or tf if pt returns a torch data tensordataset if tf returns a tf data dataset threads multiple processing threads returns list of data processors squad squadfeatures example python processor squadv2processor examples processor get_dev_examples data_dir features squad_convert_examples_to_features examples examples tokenizer tokenizer max_seq_length args max_seq_length doc_stride args doc_stride max_query_length args max_query_length is_training not evaluate defining helper methods convert to tensors and build dataset why have we split the batch into a tuple pytorch just has a list of tensors processor for the squad data set overridden by squadv1processor and squadv2processor used by the version 1 1 and version 2 0 of squad respectively creates a list of data processors squad squadexample using a tfds dataset args dataset the tfds dataset loaded from tensorflow_datasets load squad evaluate boolean specifying if in evaluation mode or in training mode returns list of squadexample examples python import tensorflow_datasets as tfds dataset tfds load squad training_examples get_examples_from_dataset dataset evaluate false evaluation_examples get_examples_from_dataset dataset evaluate true returns the training examples from the data directory args data_dir directory containing the data files used for training and evaluating filename none by default specify this if the training file has a different name than the original one which is train v1 1 json and train v2 0 json for squad versions 1 1 and 2 0 respectively returns the evaluation example from the data directory args data_dir directory containing the data files used for training and evaluating filename none by default specify this if the evaluation file has a different name than the original one which is dev v1 1 json and dev v2 0 json for squad versions 1 1 and 2 0 respectively a single training test example for the squad dataset as loaded from disk args qas_id the example s unique identifier question_text the question string context_text the context string answer_text the answer string start_position_character the character position of the start of the answer title the title of the example answers none by default this is used during evaluation holds answers as well as their start positions is_impossible false by default set to true if the example has no possible answer split on whitespace so that different tokens may be attributed to their original position start and end positions only has a value during evaluation single squad example features to be fed to a model those features are model specific and can be crafted from data processors squad squadexample using the method transformers data processors squad squad_convert_examples_to_features method args input_ids indices of input sequence tokens in the vocabulary attention_mask mask to avoid performing attention on padding token indices token_type_ids segment token indices to indicate first and second portions of the inputs cls_index the index of the cls token p_mask mask identifying tokens that can be answers vs tokens that cannot mask with 1 for tokens than cannot be in the answer and 0 for token that can be in an answer example_index the index of the example unique_id the unique feature identifier paragraph_len the length of the context token_is_max_context list of booleans identifying which tokens have their maximum context in this feature object if a token does not have their maximum context in this feature object it means that another feature object has more information related to that token and should be prioritized over this feature for that token tokens list of tokens corresponding to the input ids token_to_orig_map mapping between the tokens and the original text needed in order to identify the answer start_position start of the answer token index end_position end of the answer token index encoding optionally store the batchencoding with the fast tokenizer alignment methods constructs a squadresult which can be used to evaluate a model s output on the squad dataset args unique_id the unique identifier corresponding to that example start_logits the logits corresponding to the start of the answer end_logits the logits corresponding to the end of the answer
import json
import os
from functools import partial
from multiprocessing import Pool, cpu_count

import numpy as np
from tqdm import tqdm

from ...models.bert.tokenization_bert import whitespace_tokenize
from ...tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase, TruncationStrategy
from ...utils import is_tf_available, is_torch_available, logging
from .utils import DataProcessor



MULTI_SEP_TOKENS_TOKENIZERS_SET = {"roberta", "camembert", "bart", "mpnet"}


if is_torch_available():
    import torch
    from torch.utils.data import TensorDataset

if is_tf_available():
    import tensorflow as tf

logger = logging.get_logger(__name__)


def _improve_answer_span(doc_tokens, input_start, input_end, tokenizer, orig_answer_text):
    
    tok_answer_text = " ".join(tokenizer.tokenize(orig_answer_text))

    for new_start in range(input_start, input_end + 1):
        for new_end in range(input_end, new_start - 1, -1):
            text_span = " ".join(doc_tokens[new_start : (new_end + 1)])
            if text_span == tok_answer_text:
                return (new_start, new_end)

    return (input_start, input_end)


def _check_is_max_context(doc_spans, cur_span_index, position):
    
    best_score = None
    best_span_index = None
    for span_index, doc_span in enumerate(doc_spans):
        end = doc_span.start + doc_span.length - 1
        if position < doc_span.start:
            continue
        if position > end:
            continue
        num_left_context = position - doc_span.start
        num_right_context = end - position
        score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
        if best_score is None or score > best_score:
            best_score = score
            best_span_index = span_index

    return cur_span_index == best_span_index


def _new_check_is_max_context(doc_spans, cur_span_index, position):
    
    
    
    best_score = None
    best_span_index = None
    for span_index, doc_span in enumerate(doc_spans):
        end = doc_span["start"] + doc_span["length"] - 1
        if position < doc_span["start"]:
            continue
        if position > end:
            continue
        num_left_context = position - doc_span["start"]
        num_right_context = end - position
        score = min(num_left_context, num_right_context) + 0.01 * doc_span["length"]
        if best_score is None or score > best_score:
            best_score = score
            best_span_index = span_index

    return cur_span_index == best_span_index


def _is_whitespace(c):
    if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
        return True
    return False


def squad_convert_example_to_features(
    example, max_seq_length, doc_stride, max_query_length, padding_strategy, is_training
):
    features = []
    if is_training and not example.is_impossible:
        
        start_position = example.start_position
        end_position = example.end_position

        
        actual_text = " ".join(example.doc_tokens[start_position : (end_position + 1)])
        cleaned_answer_text = " ".join(whitespace_tokenize(example.answer_text))
        if actual_text.find(cleaned_answer_text) == -1:
            logger.warning(f"Could not find answer: '{actual_text}' vs. '{cleaned_answer_text}'")
            return []

    tok_to_orig_index = []
    orig_to_tok_index = []
    all_doc_tokens = []
    for i, token in enumerate(example.doc_tokens):
        orig_to_tok_index.append(len(all_doc_tokens))
        if tokenizer.__class__.__name__ in [
            "RobertaTokenizer",
            "LongformerTokenizer",
            "BartTokenizer",
            "RobertaTokenizerFast",
            "LongformerTokenizerFast",
            "BartTokenizerFast",
        ]:
            sub_tokens = tokenizer.tokenize(token, add_prefix_space=True)
        else:
            sub_tokens = tokenizer.tokenize(token)
        for sub_token in sub_tokens:
            tok_to_orig_index.append(i)
            all_doc_tokens.append(sub_token)

    if is_training and not example.is_impossible:
        tok_start_position = orig_to_tok_index[example.start_position]
        if example.end_position < len(example.doc_tokens) - 1:
            tok_end_position = orig_to_tok_index[example.end_position + 1] - 1
        else:
            tok_end_position = len(all_doc_tokens) - 1

        (tok_start_position, tok_end_position) = _improve_answer_span(
            all_doc_tokens, tok_start_position, tok_end_position, tokenizer, example.answer_text
        )

    spans = []

    truncated_query = tokenizer.encode(
        example.question_text, add_special_tokens=False, truncation=True, max_length=max_query_length
    )

    
    
    tokenizer_type = type(tokenizer).__name__.replace("Tokenizer", "").lower()
    sequence_added_tokens = (
        tokenizer.model_max_length - tokenizer.max_len_single_sentence + 1
        if tokenizer_type in MULTI_SEP_TOKENS_TOKENIZERS_SET
        else tokenizer.model_max_length - tokenizer.max_len_single_sentence
    )
    sequence_pair_added_tokens = tokenizer.model_max_length - tokenizer.max_len_sentences_pair

    span_doc_tokens = all_doc_tokens
    while len(spans) * doc_stride < len(all_doc_tokens):
        
        if tokenizer.padding_side == "right":
            texts = truncated_query
            pairs = span_doc_tokens
            truncation = TruncationStrategy.ONLY_SECOND.value
        else:
            texts = span_doc_tokens
            pairs = truncated_query
            truncation = TruncationStrategy.ONLY_FIRST.value

        encoded_dict = tokenizer.encode_plus(  
            texts,
            pairs,
            truncation=truncation,
            padding=padding_strategy,
            max_length=max_seq_length,
            return_overflowing_tokens=True,
            stride=max_seq_length - doc_stride - len(truncated_query) - sequence_pair_added_tokens,
            return_token_type_ids=True,
        )

        paragraph_len = min(
            len(all_doc_tokens) - len(spans) * doc_stride,
            max_seq_length - len(truncated_query) - sequence_pair_added_tokens,
        )

        if tokenizer.pad_token_id in encoded_dict["input_ids"]:
            if tokenizer.padding_side == "right":
                non_padded_ids = encoded_dict["input_ids"][: encoded_dict["input_ids"].index(tokenizer.pad_token_id)]
            else:
                last_padding_id_position = (
                    len(encoded_dict["input_ids"]) - 1 - encoded_dict["input_ids"][::-1].index(tokenizer.pad_token_id)
                )
                non_padded_ids = encoded_dict["input_ids"][last_padding_id_position + 1 :]

        else:
            non_padded_ids = encoded_dict["input_ids"]

        tokens = tokenizer.convert_ids_to_tokens(non_padded_ids)

        token_to_orig_map = {}
        for i in range(paragraph_len):
            index = len(truncated_query) + sequence_added_tokens + i if tokenizer.padding_side == "right" else i
            token_to_orig_map[index] = tok_to_orig_index[len(spans) * doc_stride + i]

        encoded_dict["paragraph_len"] = paragraph_len
        encoded_dict["tokens"] = tokens
        encoded_dict["token_to_orig_map"] = token_to_orig_map
        encoded_dict["truncated_query_with_special_tokens_length"] = len(truncated_query) + sequence_added_tokens
        encoded_dict["token_is_max_context"] = {}
        encoded_dict["start"] = len(spans) * doc_stride
        encoded_dict["length"] = paragraph_len

        spans.append(encoded_dict)

        if "overflowing_tokens" not in encoded_dict or (
            "overflowing_tokens" in encoded_dict and len(encoded_dict["overflowing_tokens"]) == 0
        ):
            break
        span_doc_tokens = encoded_dict["overflowing_tokens"]

    for doc_span_index in range(len(spans)):
        for j in range(spans[doc_span_index]["paragraph_len"]):
            is_max_context = _new_check_is_max_context(spans, doc_span_index, doc_span_index * doc_stride + j)
            index = (
                j
                if tokenizer.padding_side == "left"
                else spans[doc_span_index]["truncated_query_with_special_tokens_length"] + j
            )
            spans[doc_span_index]["token_is_max_context"][index] = is_max_context

    for span in spans:
        
        cls_index = span["input_ids"].index(tokenizer.cls_token_id)

        
        
        p_mask = np.ones_like(span["token_type_ids"])
        if tokenizer.padding_side == "right":
            p_mask[len(truncated_query) + sequence_added_tokens :] = 0
        else:
            p_mask[-len(span["tokens"]) : -(len(truncated_query) + sequence_added_tokens)] = 0

        pad_token_indices = np.where(span["input_ids"] == tokenizer.pad_token_id)
        special_token_indices = np.asarray(
            tokenizer.get_special_tokens_mask(span["input_ids"], already_has_special_tokens=True)
        ).nonzero()

        p_mask[pad_token_indices] = 1
        p_mask[special_token_indices] = 1

        
        p_mask[cls_index] = 0

        span_is_impossible = example.is_impossible
        start_position = 0
        end_position = 0
        if is_training and not span_is_impossible:
            
            
            doc_start = span["start"]
            doc_end = span["start"] + span["length"] - 1
            out_of_span = False

            if not (tok_start_position >= doc_start and tok_end_position <= doc_end):
                out_of_span = True

            if out_of_span:
                start_position = cls_index
                end_position = cls_index
                span_is_impossible = True
            else:
                if tokenizer.padding_side == "left":
                    doc_offset = 0
                else:
                    doc_offset = len(truncated_query) + sequence_added_tokens

                start_position = tok_start_position - doc_start + doc_offset
                end_position = tok_end_position - doc_start + doc_offset

        features.append(
            SquadFeatures(
                span["input_ids"],
                span["attention_mask"],
                span["token_type_ids"],
                cls_index,
                p_mask.tolist(),
                example_index=0,  
                unique_id=0,
                paragraph_len=span["paragraph_len"],
                token_is_max_context=span["token_is_max_context"],
                tokens=span["tokens"],
                token_to_orig_map=span["token_to_orig_map"],
                start_position=start_position,
                end_position=end_position,
                is_impossible=span_is_impossible,
                qas_id=example.qas_id,
            )
        )
    return features


def squad_convert_example_to_features_init(tokenizer_for_convert: PreTrainedTokenizerBase):
    global tokenizer
    tokenizer = tokenizer_for_convert


def squad_convert_examples_to_features(
    examples,
    tokenizer,
    max_seq_length,
    doc_stride,
    max_query_length,
    is_training,
    padding_strategy="max_length",
    return_dataset=False,
    threads=1,
    tqdm_enabled=True,
):
    
    
    features = []

    threads = min(threads, cpu_count())
    with Pool(threads, initializer=squad_convert_example_to_features_init, initargs=(tokenizer,)) as p:
        annotate_ = partial(
            squad_convert_example_to_features,
            max_seq_length=max_seq_length,
            doc_stride=doc_stride,
            max_query_length=max_query_length,
            padding_strategy=padding_strategy,
            is_training=is_training,
        )
        features = list(
            tqdm(
                p.imap(annotate_, examples, chunksize=32),
                total=len(examples),
                desc="convert squad examples to features",
                disable=not tqdm_enabled,
            )
        )

    new_features = []
    unique_id = 1000000000
    example_index = 0
    for example_features in tqdm(
        features, total=len(features), desc="add example index and unique id", disable=not tqdm_enabled
    ):
        if not example_features:
            continue
        for example_feature in example_features:
            example_feature.example_index = example_index
            example_feature.unique_id = unique_id
            new_features.append(example_feature)
            unique_id += 1
        example_index += 1
    features = new_features
    del new_features
    if return_dataset == "pt":
        if not is_torch_available():
            raise RuntimeError("PyTorch must be installed to return a PyTorch dataset.")

        
        all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
        all_attention_masks = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
        all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
        all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)
        all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)
        all_is_impossible = torch.tensor([f.is_impossible for f in features], dtype=torch.float)

        if not is_training:
            all_feature_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
            dataset = TensorDataset(
                all_input_ids, all_attention_masks, all_token_type_ids, all_feature_index, all_cls_index, all_p_mask
            )
        else:
            all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)
            all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)
            dataset = TensorDataset(
                all_input_ids,
                all_attention_masks,
                all_token_type_ids,
                all_start_positions,
                all_end_positions,
                all_cls_index,
                all_p_mask,
                all_is_impossible,
            )

        return features, dataset
    elif return_dataset == "tf":
        if not is_tf_available():
            raise RuntimeError("TensorFlow must be installed to return a TensorFlow dataset.")

        def gen():
            for i, ex in enumerate(features):
                if ex.token_type_ids is None:
                    yield (
                        {
                            "input_ids": ex.input_ids,
                            "attention_mask": ex.attention_mask,
                            "feature_index": i,
                            "qas_id": ex.qas_id,
                        },
                        {
                            "start_positions": ex.start_position,
                            "end_positions": ex.end_position,
                            "cls_index": ex.cls_index,
                            "p_mask": ex.p_mask,
                            "is_impossible": ex.is_impossible,
                        },
                    )
                else:
                    yield (
                        {
                            "input_ids": ex.input_ids,
                            "attention_mask": ex.attention_mask,
                            "token_type_ids": ex.token_type_ids,
                            "feature_index": i,
                            "qas_id": ex.qas_id,
                        },
                        {
                            "start_positions": ex.start_position,
                            "end_positions": ex.end_position,
                            "cls_index": ex.cls_index,
                            "p_mask": ex.p_mask,
                            "is_impossible": ex.is_impossible,
                        },
                    )

        
        if "token_type_ids" in tokenizer.model_input_names:
            train_types = (
                {
                    "input_ids": tf.int32,
                    "attention_mask": tf.int32,
                    "token_type_ids": tf.int32,
                    "feature_index": tf.int64,
                    "qas_id": tf.string,
                },
                {
                    "start_positions": tf.int64,
                    "end_positions": tf.int64,
                    "cls_index": tf.int64,
                    "p_mask": tf.int32,
                    "is_impossible": tf.int32,
                },
            )

            train_shapes = (
                {
                    "input_ids": tf.TensorShape([None]),
                    "attention_mask": tf.TensorShape([None]),
                    "token_type_ids": tf.TensorShape([None]),
                    "feature_index": tf.TensorShape([]),
                    "qas_id": tf.TensorShape([]),
                },
                {
                    "start_positions": tf.TensorShape([]),
                    "end_positions": tf.TensorShape([]),
                    "cls_index": tf.TensorShape([]),
                    "p_mask": tf.TensorShape([None]),
                    "is_impossible": tf.TensorShape([]),
                },
            )
        else:
            train_types = (
                {"input_ids": tf.int32, "attention_mask": tf.int32, "feature_index": tf.int64, "qas_id": tf.string},
                {
                    "start_positions": tf.int64,
                    "end_positions": tf.int64,
                    "cls_index": tf.int64,
                    "p_mask": tf.int32,
                    "is_impossible": tf.int32,
                },
            )

            train_shapes = (
                {
                    "input_ids": tf.TensorShape([None]),
                    "attention_mask": tf.TensorShape([None]),
                    "feature_index": tf.TensorShape([]),
                    "qas_id": tf.TensorShape([]),
                },
                {
                    "start_positions": tf.TensorShape([]),
                    "end_positions": tf.TensorShape([]),
                    "cls_index": tf.TensorShape([]),
                    "p_mask": tf.TensorShape([None]),
                    "is_impossible": tf.TensorShape([]),
                },
            )

        return tf.data.Dataset.from_generator(gen, train_types, train_shapes)
    else:
        return features


class SquadProcessor(DataProcessor):
    

    train_file = None
    dev_file = None

    def _get_example_from_tensor_dict(self, tensor_dict, evaluate=False):
        if not evaluate:
            answer = tensor_dict["answers"]["text"][0].numpy().decode("utf-8")
            answer_start = tensor_dict["answers"]["answer_start"][0].numpy()
            answers = []
        else:
            answers = [
                {"answer_start": start.numpy(), "text": text.numpy().decode("utf-8")}
                for start, text in zip(tensor_dict["answers"]["answer_start"], tensor_dict["answers"]["text"])
            ]

            answer = None
            answer_start = None

        return SquadExample(
            qas_id=tensor_dict["id"].numpy().decode("utf-8"),
            question_text=tensor_dict["question"].numpy().decode("utf-8"),
            context_text=tensor_dict["context"].numpy().decode("utf-8"),
            answer_text=answer,
            start_position_character=answer_start,
            title=tensor_dict["title"].numpy().decode("utf-8"),
            answers=answers,
        )

    def get_examples_from_dataset(self, dataset, evaluate=False):
        

        if evaluate:
            dataset = dataset["validation"]
        else:
            dataset = dataset["train"]

        examples = []
        for tensor_dict in tqdm(dataset):
            examples.append(self._get_example_from_tensor_dict(tensor_dict, evaluate=evaluate))

        return examples

    def get_train_examples(self, data_dir, filename=None):
        
        if data_dir is None:
            data_dir = ""

        if self.train_file is None:
            raise ValueError("SquadProcessor should be instantiated via SquadV1Processor or SquadV2Processor")

        with open(
            os.path.join(data_dir, self.train_file if filename is None else filename), "r", encoding="utf-8"
        ) as reader:
            input_data = json.load(reader)["data"]
        return self._create_examples(input_data, "train")

    def get_dev_examples(self, data_dir, filename=None):
        
        if data_dir is None:
            data_dir = ""

        if self.dev_file is None:
            raise ValueError("SquadProcessor should be instantiated via SquadV1Processor or SquadV2Processor")

        with open(
            os.path.join(data_dir, self.dev_file if filename is None else filename), "r", encoding="utf-8"
        ) as reader:
            input_data = json.load(reader)["data"]
        return self._create_examples(input_data, "dev")

    def _create_examples(self, input_data, set_type):
        is_training = set_type == "train"
        examples = []
        for entry in tqdm(input_data):
            title = entry["title"]
            for paragraph in entry["paragraphs"]:
                context_text = paragraph["context"]
                for qa in paragraph["qas"]:
                    qas_id = qa["id"]
                    question_text = qa["question"]
                    start_position_character = None
                    answer_text = None
                    answers = []

                    is_impossible = qa.get("is_impossible", False)
                    if not is_impossible:
                        if is_training:
                            answer = qa["answers"][0]
                            answer_text = answer["text"]
                            start_position_character = answer["answer_start"]
                        else:
                            answers = qa["answers"]

                    example = SquadExample(
                        qas_id=qas_id,
                        question_text=question_text,
                        context_text=context_text,
                        answer_text=answer_text,
                        start_position_character=start_position_character,
                        title=title,
                        is_impossible=is_impossible,
                        answers=answers,
                    )
                    examples.append(example)
        return examples


class SquadV1Processor(SquadProcessor):
    train_file = "train-v1.1.json"
    dev_file = "dev-v1.1.json"


class SquadV2Processor(SquadProcessor):
    train_file = "train-v2.0.json"
    dev_file = "dev-v2.0.json"


class SquadExample:
    

    def __init__(
        self,
        qas_id,
        question_text,
        context_text,
        answer_text,
        start_position_character,
        title,
        answers=[],
        is_impossible=False,
    ):
        self.qas_id = qas_id
        self.question_text = question_text
        self.context_text = context_text
        self.answer_text = answer_text
        self.title = title
        self.is_impossible = is_impossible
        self.answers = answers

        self.start_position, self.end_position = 0, 0

        doc_tokens = []
        char_to_word_offset = []
        prev_is_whitespace = True

        
        for c in self.context_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)

        self.doc_tokens = doc_tokens
        self.char_to_word_offset = char_to_word_offset

        
        if start_position_character is not None and not is_impossible:
            self.start_position = char_to_word_offset[start_position_character]
            self.end_position = char_to_word_offset[
                min(start_position_character + len(answer_text) - 1, len(char_to_word_offset) - 1)
            ]


class SquadFeatures:
    

    def __init__(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        cls_index,
        p_mask,
        example_index,
        unique_id,
        paragraph_len,
        token_is_max_context,
        tokens,
        token_to_orig_map,
        start_position,
        end_position,
        is_impossible,
        qas_id: str = None,
        encoding: BatchEncoding = None,
    ):
        self.input_ids = input_ids
        self.attention_mask = attention_mask
        self.token_type_ids = token_type_ids
        self.cls_index = cls_index
        self.p_mask = p_mask

        self.example_index = example_index
        self.unique_id = unique_id
        self.paragraph_len = paragraph_len
        self.token_is_max_context = token_is_max_context
        self.tokens = tokens
        self.token_to_orig_map = token_to_orig_map

        self.start_position = start_position
        self.end_position = end_position
        self.is_impossible = is_impossible
        self.qas_id = qas_id

        self.encoding = encoding


class SquadResult:
    

    def __init__(self, unique_id, start_logits, end_logits, start_top_index=None, end_top_index=None, cls_logits=None):
        self.start_logits = start_logits
        self.end_logits = end_logits
        self.unique_id = unique_id

        if start_top_index:
            self.start_top_index = start_top_index
            self.end_top_index = end_top_index
            self.cls_logits = cls_logits