U 4Af@s\ddlZddlZddlmZddlmZmZddlZddl m Z ddl m Z ddl m Z mZmZddlmZmZmZd d lmZd d d dhZerddlZddlmZerddlZeeZddZddZddZ ddZ!ddZ"edddZ#d.d d!Z$Gd"d#d#eZ%Gd$d%d%e%Z&Gd&d'd'e%Z'Gd(d)d)Z(Gd*d+d+Z)Gd,d-d-Z*dS)/N)partial)Pool cpu_count)tqdm)whitespace_tokenize) BatchEncodingPreTrainedTokenizerBaseTruncationStrategy)is_tf_availableis_torch_availablelogging) DataProcessorZrobertaZ camembertZbartZmpnet) TensorDatasetc Cspd||}t||dD]H}t||ddD]2}d|||d}||kr2||fSq2q||fS)zFReturns tokenized answer spans that better match the annotated answer. r)jointokenizerange) doc_tokensZ input_startZ input_end tokenizerZorig_answer_textZtok_answer_textZ new_startZnew_endZ text_spanrF/tmp/pip-unpacked-wheel-zw5xktn0/transformers/data/processors/squad.py_improve_answer_span+src Csd}d}t|D]l\}}|j|jd}||jkr4q||kr>q||j}||} t|| d|j} |dkst| |kr| }|}q||kS):Check if this is the 'max context' doc span for the token.Nr{Gz?) enumeratestartlengthmin Z doc_spansZcur_span_indexpositionZ best_scoreZbest_span_indexZ span_indexZdoc_spanendZnum_left_contextZnum_right_contextZscorerrr_check_is_max_context8s  r$c Csd}d}t|D]v\}}|d|dd}||dkr:q||krDq||d}||} t|| d|d} |dks~| |kr| }|}q||kS)rNrrrr)rr r!rrr_new_check_is_max_contextLs  r%cCs4|dks,|dks,|dks,|dks,t|dkr0dSdS)Nr   i/ TF)ord)crrr_is_whitespacebs,r+c1Csg}|rl|jsl|j}|j}d|j||d} dt|j} | | dkrlt d| d| dgSg} g} g} t |jD]Z\}}| t | t jjdkrt j|dd }n t |}|D]}| || |qq|rB|jsB| |j}|jt |jdkr | |jdd}n t | d}t| ||t |j\}}g}t j|jd d|d }tt jd d }|tkrt jt jdn t jt j}t jt j}| }t ||t | krt jdkr|}|}tjj}n|}|}tjj}t j |||||d||t ||dd}t!t | t |||t ||}t j"|dkrt jdkrt|dd|d#t j"}n>t |dd|dddd#t j"} |d| dd}n|d}t $|}!i}"t%|D]>}t jdkrt |||n|}#| t ||||"|#<q||d<|!|d<|"|d<t |||d<i|d<t |||d<||d<| |d|ksd|krt |ddkrq|d}qt%t |D]b}$t%||$dD]J}%t&||$|$||%}&t jdkr|%n||$d|%}#|&||$d|#<qq|D]}'|'d#t j'}(t()|'d})t jdkrTd|)t ||d<n d|)t |'d t || <t(*|'dt j"k}*t(+t j,|'ddd-}+d|)|*<d|)|+<d|)|(<|j},d}d}|rX|,sX|'d}-|'d|'dd}.d }/||-kr ||.ksd}/|/r"|(}|(}d},n6t jdkr4d}0n t ||}0||-|0}||-|0}| t.|'d|'d|'d|(|)/dd|'d|'d|'d|'d|||,|j0dq|S)NrrrzCould not find answer: 'z' vs. '')ZRobertaTokenizerZLongformerTokenizerZ BartTokenizerZRobertaTokenizerFastZLongformerTokenizerFastZBartTokenizerFastT)Zadd_prefix_spaceF)Zadd_special_tokens truncation max_length Tokenizerright)r-paddingr.Zreturn_overflowing_tokensZstrideZreturn_token_type_ids input_ids paragraph_lentokenstoken_to_orig_mapZ*truncated_query_with_special_tokens_lengthtoken_is_max_contextrrZoverflowing_tokensrlefttoken_type_ids)Zalready_has_special_tokensattention_mask) example_index unique_idr4r7r5r6start_position end_position is_impossibleqas_id)1r?r=r>rrr answer_textfindloggerwarningrappendlenr __class____name__rrencode question_texttypereplacelowerMULTI_SEP_TOKENS_TOKENIZERS_SETZmodel_max_lengthZmax_len_single_sentenceZmax_len_sentences_pairZ padding_sider Z ONLY_SECONDvalueZ ONLY_FIRSTZ encode_plusr Z pad_token_idindexZconvert_ids_to_tokensrr%Z cls_token_idnpZ ones_likewhereZasarrayZget_special_tokens_maskZnonzero SquadFeaturestolistr@)1examplemax_seq_length doc_stridemax_query_lengthpadding_strategy is_trainingfeaturesr=r>Z actual_textZcleaned_answer_textZtok_to_orig_indexZorig_to_tok_indexZall_doc_tokensitokenZ sub_tokensZ sub_tokenZtok_start_positionZtok_end_positionZspansZtruncated_queryZtokenizer_typeZsequence_added_tokensZsequence_pair_added_tokensZspan_doc_tokensZtextspairsr-Z encoded_dictr4Znon_padded_idsZlast_padding_id_positionr5r6rPZdoc_span_indexjZis_max_contextspan cls_indexp_maskZpad_token_indicesZspecial_token_indicesZspan_is_impossibleZ doc_startZdoc_endZ out_of_spanZ doc_offsetrrr!squad_convert_example_to_featureshs4            (              rcZtokenizer_for_convertcCs|adSN)rrdrrr&squad_convert_example_to_features_init7srfr.FTc  spgt|t}t|t|fd@} tt|||||d} tt| j| |ddt |d| dW5QRXg} d} d}tt d | dD]:}|sq|D]"}||_ | |_ | || d 7} q|d 7}q| ~ |d krt std tjd dDtjd}tjddDtjd}tjddDtjd}tjddDtjd}tjddDtjd}tjddDtjd}|stj|dtjd}t||||||}nJtjddDtjd}tjddDtjd}t||||||||}|fS|dkrhtstdfdd}d|jkrtjtjtjtjtjdtjtjtjtjtjdf}tdgtdgtdgtgtgdtgtgtgtdgtgdf}ntjtjtjtjdtjtjtjtjtjdf}tdgtdgtgtgdtgtgtgtdgtgdf}tjj|||SSdS) a 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, ) ```)Z initializerZinitargs)rVrWrXrYrZ ) chunksizez"convert squad examples to features)totaldescdisableiʚ;rzadd example index and unique idrptz6PyTorch must be installed to return a PyTorch dataset.cSsg|] }|jqSr)r3.0frrr sz6squad_convert_examples_to_features..)ZdtypecSsg|] }|jqSr)r:rmrrrrpscSsg|] }|jqSr)r9rmrrrrpscSsg|] }|jqSr)rarmrrrrpscSsg|] }|jqSr)rbrmrrrrpscSsg|] }|jqSr)r?rmrrrrpscSsg|] }|jqSr)r=rmrrrrpscSsg|] }|jqSr)r>rmrrrrpstfzrarbr?)r\exr[rrgens6 z/squad_convert_examples_to_features..genr9rurtNrr) r rrrfrrclistrimaprFr;r<rEr RuntimeErrortorchZtensorlongfloatZarangesizerr Zmodel_input_namesrqZint32Zint64stringZ TensorShapedataZDatasetZfrom_generator)examplesrrVrWrXrZrYZreturn_datasetthreadsZ tqdm_enabledpZ annotate_Z new_featuresr<r;Zexample_featuresZexample_featureZ all_input_idsZall_attention_masksZall_token_type_idsZ all_cls_indexZ all_p_maskZall_is_impossibleZall_feature_indexdatasetZall_start_positionsZall_end_positionsrxZ train_typesZ train_shapesrrwr"squad_convert_examples_to_features<s0          %        rc@sHeZdZdZdZdZdddZdddZddd Zdd d Z d d Z dS)SquadProcessorz Processor for the SQuAD data set. overridden by SquadV1Processor and SquadV2Processor, used by the version 1.1 and version 2.0 of SQuAD, respectively. NFc Cs|s8|dddd}|ddd}g}n,ddt|dd|ddD}d}d}t|dd|d d|d d|||d d|d S) Nanswerstextrutf-8 answer_startcSs(g|] \}}||ddqS)r)rr)numpydecode)rnrrrrrrp,sz@SquadProcessor._get_example_from_tensor_dict..idquestioncontexttitle)r@rJ context_textrAstart_position_characterrr)rrzip SquadExample)self tensor_dictevaluateanswerrrrrr_get_example_from_tensor_dict&s$z,SquadProcessor._get_example_from_tensor_dictcCs@|r|d}n|d}g}t|D]}||j||dq"|S)av 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) ```Z validationtrain)r)rrEr)rrrrrrrrget_examples_from_dataset>s  z(SquadProcessor.get_examples_from_datasetc Csj|dkr d}|jdkrtdttj||dkr6|jn|ddd}t|d}W5QRX||dS) a 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. Nr0NSquadProcessor should be instantiated via SquadV1Processor or SquadV2Processorrrencodingrr) train_file ValueErroropenospathrjsonload_create_examplesrdata_dirfilenamereader input_datarrrget_train_examples_s  z!SquadProcessor.get_train_examplesc Csj|dkr d}|jdkrtdttj||dkr6|jn|ddd}t|d}W5QRX||dS) a 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. Nr0rrrrrdev) dev_filerrrrrrrrrrrrget_dev_examplesus  zSquadProcessor.get_dev_examplesc Cs|dk}g}t|D]}|d}|dD]}|d}|dD]|} | d} | d} d} d} g}| dd }|s|r| d d }|d } |d } n| d }t| | || | |||d}||q       zSquadProcessor._create_examples)F)F)N)N) rH __module__ __qualname____doc__rrrrrrrrrrrrs  !  rc@seZdZdZdZdS)SquadV1Processorztrain-v1.1.jsonz dev-v1.1.jsonNrHrrrrrrrrrsrc@seZdZdZdZdS)SquadV2Processorztrain-v2.0.jsonz dev-v2.0.jsonNrrrrrrsrc@seZdZdZgdfddZdS)raT 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. Fc Cs||_||_||_||_||_||_||_d\|_|_g} g} d} |jD]H} t | rZd} n$| rj| | n| d| 7<d} | t | dqH| |_ | |_ |dk r|s| ||_| t|t |dt | d|_dS)N)rrTrFr)r@rJrrArr?rr=r>r+rErFrchar_to_word_offsetr ) rr@rJrrArrrr?rrZprev_is_whitespacer*rrr__init__s4      zSquadExample.__init__NrHrrrrrrrrrsrc@s"eZdZdZdeedddZdS)rSa 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. N)r@rcCsd||_||_||_||_||_||_||_||_| |_| |_ | |_ | |_ | |_ ||_ ||_||_dSre)r3r:r9rarbr;r<r4r7r5r6r=r>r?r@r)rr3r:r9rarbr;r<r4r7r5r6r=r>r?r@rrrrrs zSquadFeatures.__init__)NN)rHrrrstrrrrrrrrSs*rSc@seZdZdZdddZdS) SquadResultaJ 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 NcCs,||_||_||_|r(||_||_||_dSre) start_logits end_logitsr<start_top_index end_top_index cls_logits)rr<rrrrrrrrrEszSquadResult.__init__)NNNrrrrrr;s r)r.FrT)+rr functoolsrmultiprocessingrrrrQrZmodels.bert.tokenization_bertrZtokenization_utils_baserr r utilsr r r rrNr|Ztorch.utils.datarZ tensorflowrqZ get_loggerrHrCrr$r%r+rcrfrrrrrrSrrrrrsH        P  b?C