Aluode/PerceptionLabPortable
0
1o
2 .�Yi�� � @ sN d dl Zd dlZd dlZd dlmZ d dlmZ d dlmZ d dl m3Z4m Z m Z m
Z
d dlZddlmZ ddlmZ e5Z e ee geee6f f ZG d d7� d8�Zd d � Zd7dee deee9f fdd�ZeG dd� de��Z dee deee10f fdd�Z dee deee11f fdd�Z dee deee12f fdd�Z eG dd� d��Z eG d d � d e��Z!d8d e e" fd d �Z#d8d e e" fd!d"�Z$d8d e e" fd#d$�Z%eG d%d&� d&e��Z&eG d'd(� d(��Z'eG d)d*� d*e��Z(eG d+d,� d,e(��Z)dee13 fd-d.�Z*dej+e14 fd/d0�Z,eG d1d2� d2e(��Z-eG d3d4� d4e��Z.eG d5d6� d6e ��Z/dS )9� N)�Mapping)� dataclass)�randint)�Any�Callable�Optional�Union� )�PreTrainedTokenizerBase)�PaddingStrategyc @ s e Zd Zddee fdd�ZdS )�DataCollatorMixinN�return_tensorsc C sT |d u r| j }|dkr| �|�S |dkr| �|�S |dkr"| �|�S td|� d���)N�tf�pt�npz Framework 'z' not recognized!)r
�tf_call�15torch_call�16numpy_call�17ValueError��self�featuresr
� r ��E:\DocsHouse\542 percep lab latest\PerceptionLab\PerceptionLab_Portable\python_embed\Lib\site-packages\transformers/data/data_collator.py�__call__' s 181920zDataCollatorMixin.__call__�N)�__name__�21__module__� __qualname__r �strr r r r r r & s r c O s` t | d�s
| j|i |��S | j�dd�}d| jd< z| j|i |��}W || jd< |S || jd< w )zz22 Pads without triggering the warning about how using the pad function is sub-optimal when using a fast tokenizer.23 �deprecation_warningsz Asking-to-pad-a-fast-tokenizerFT)�hasattr�padr �get)� tokenizerZpad_argsZ24pad_kwargsZ
warning_state�paddedr r r �"pad_without_fast_tokenizer_warning4 s 252627 �r& r r �returnc C s4 |dkrt | �S |dkrt| �S |dkrt| �S dS )a� 28 Very simple data collator that simply collates batches of dict-like objects and performs special handling for29 potential keys named:30 31 - `label`: handles a single value (int or float) per object32 - `label_ids`: handles a list of values per object33 34 Does not do any additional preprocessing: property names of the input object will be used as corresponding inputs35 to the model. See glue and ner for example of how it's useful.36 r r r N)�torch_default_data_collator�tf_default_data_collator�numpy_default_data_collator)r r
r r r �default_data_collatorJ s �r+ c @ sF e Zd ZU dZdZeed< d deeee f deee f fdd�Z37dS )38�DefaultDataCollatora* 39 Very simple data collator that simply collates batches of dict-like objects and performs special handling for40 potential keys named:41 42 - `label`: handles a single value (int or float) per object43 - `label_ids`: handles a list of values per object44 45 Does not do any additional preprocessing: property names of the input object will be used as corresponding inputs46 to the model. See glue and ner for example of how it's useful.47 48 This is an object (like other data collators) rather than a pure function like default_data_collator. This can be49 helpful if you need to set a return_tensors value at initialization.50 51 Args:52 return_tensors (`str`, *optional*, defaults to `"pt"`):53 The type of Tensor to return. Allowable values are "np", "pt" and "tf".54 r r
Nr r' c C s |d u r| j }t||�S r )r
r+ r r r r r y s 55z DefaultDataCollator.__call__r ) r r r �__doc__r
r �__annotations__�list�dictr r r r r r r, c s 56 ,r, c s� dd l }t| d t�sdd� | D �} | d }i }d|v rN|d d urNt|d |j�r0|d �� n|d }t|t�r<|jn|j}|jdd� | D �|d�|d< n<d|v r�|d d ur�t|d |j�rm|� d d� | D ��|d< n t|d d t�ry|jn|j}|jd57d� | D �|d�|d< |�58� D ]L\� }� d vr�|d ur�t|t �s�t||j�r�|� � fd d�| D ��|� < q�t|t j
�r�|�t � � fd
d�| D ���|� < q�|�� fdd�| D ��|� < q�|S )Nr c S � g | ]}t |��qS r ��vars��.0�fr r r �59<listcomp>� � z/torch_default_data_collator.<locals>.<listcomp>�labelc S � g | ]}|d �qS �r9 r r4 r r r r7 � r8 ��dtype�labels� label_idsc S r: �r? r r4 r r r r7 � r8 c S r: r@ r r4 r r r r7 � r8 �r9 r? c � g | ]}|� �qS r r r4 ��kr r r7 � r8 c rB r r r4 rC r r r7 � r8 c rB r r r4 rC r r r7 � r8 )�torch�60isinstancer �Tensor�item�int�long�float�tensor�stack�itemsr r �ndarrayZ61from_numpy)r rE �first�batchr9 r= �vr rC r r( s. $ $ �r( c s� dd l }t| d t�sdd� | D �} | d }i }d|v r%|d d ur%d�n d|v r2|d d ur2d�nd|v r?|d d ur?d�nd ��d ur�t|� |j�rZ|� jjrV|jn|j}nBt|� tj tj62f�rvt� |� jtj �rr|jn|j}n&t|� t
tf�r�t|� d t�r�|jn|j}n
t|� t�r�|jn|j}|j�fdd�| D �|d�|d< |�� D ]7\� }� d vr�|d ur�t|t�s�t||jtj f�r�|�� fd63d�| D ��|� < q�|�� fd d�| D ��|� < q�|S ) Nr c S r1 r r2 r4 r r r r7 � r8 z,tf_default_data_collator.<locals>.<listcomp>r9 r? r> c rB r r r4 )�label_col_namer r r7 � r8 r< )r9 r? r> c rB r r r4 rC r r r7 � r8 c rB r r r4 rC r r r7 � r8 )�64tensorflowrF r rG r= �65is_integer�int64�float32r rO ZgenericZ66issubdtype�integer�tupler/ rI �convert_to_tensorrN r rM )r r rP rQ r= rR r )rD rS r r) � s8 " �r) c s� t | d t�sdd� | D �} | d }i }d|v rJ|d d urJt |d tj�r,|d �� n|d }t |t�r8tjntj}tjdd� | D �|d�|d< n<d|v r�|d d ur�t |d tj�rit� d d� | D ��|d< n t |d d t�rutjntj}tjd67d� | D �|d�|d< |�68� D ]4\� }� d vr�|d ur�t |t �s�t |tj�r�t� � fd d�| D ��|� < q�t�� fd
d�| D ��|� < q�|S )Nr c S r1 r r2 r4 r r r r7 � r8 z/numpy_default_data_collator.<locals>.<listcomp>r9 c S r: r; r r4 r r r r7 � r8 r< r> r? c S r: r@ r r4 r r r r7 � r8 c S r: r@ r r4 r r r r7 � r8 rA c rB r r r4 rC r r r7 � r8 c rB r r r4 rC r r r7 � r8 ) rF r r rO rH rI rV rW �arrayrM rN r )r rP rQ r9 r= rR r rC r r* � s( $ �r* c @ s� e Zd ZU dZeed< dZeee e69f ed< dZ e e
ed< dZe e
ed< dZe ed < d70eee ef d ee ef fd d
�ZdS )�DataCollatorWithPaddinga� 71 Data collator that will dynamically pad the inputs received.72 73 Args:74 tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):75 The tokenizer used for encoding the data.76 padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):77 Select a strategy to pad the returned sequences (according to the model's padding side and padding index)78 among:79 80 - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single81 sequence is provided).82 - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum83 acceptable input length for the model if that argument is not provided.84 - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).85 max_length (`int`, *optional*):86 Maximum length of the returned list and optionally padding length (see above).87 pad_to_multiple_of (`int`, *optional*):88 If set will pad the sequence to a multiple of the provided value.89 90 This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=91 7.0 (Volta).92 return_tensors (`str`, *optional*, defaults to `"pt"`):93 The type of Tensor to return. Allowable values are "np", "pt" and "tf".94 r$ T�paddingN�95max_length�pad_to_multiple_ofr r
r r' c C sV t | j|| j| j| j| jd�}d|v r |d |d< |d= d|v r)|d |d< |d= |S )N�r] r^ r_ r
r9 r> r? )r& r$ r] r^ r_ r
)r r rQ r r r r s � z DataCollatorWithPadding.__call__)r r r r- r96 r. r] r �boolr r r^ r rI r_ r
r/ r0 r r r r r r r\ � s 97 *r\ c @ s� e Zd ZU dZeed< dZeee e98f ed< dZ e e
ed< dZe e
ed< dZe
ed < d99Ze ed < d d
� Zdd� Zdd� ZdS )�"DataCollatorForTokenClassificationa� 100 Data collator that will dynamically pad the inputs received, as well as the labels.101 102 Args:103 tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):104 The tokenizer used for encoding the data.105 padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):106 Select a strategy to pad the returned sequences (according to the model's padding side and padding index)107 among:108 109 - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single110 sequence is provided).111 - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum112 acceptable input length for the model if that argument is not provided.113 - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).114 max_length (`int`, *optional*):115 Maximum length of the returned list and optionally padding length (see above).116 pad_to_multiple_of (`int`, *optional*):117 If set will pad the sequence to a multiple of the provided value.118 119 This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=120 7.0 (Volta).121 label_pad_token_id (`int`, *optional*, defaults to -100):122 The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).123 return_tensors (`str`, *optional*, defaults to `"pt"`):124 The type of Tensor to return. Allowable values are "np", "pt" and "tf".125 r$ Tr] Nr^ r_ �����label_pad_token_idr r
c s� dd l �d|d v r dnd� � |d v r � fdd�|D �nd }� fdd�|D �}t�j|�j�j�jdd�}|d u r<|S |d jd126 ��jj}�fd d ��|d
kr_���fdd�|D �|� < n
���fdd�|D �|� < �j|� �j d�|� < |S )Nr r9 r> c rB r r �r5 �feature��127label_namer r r7 I r8 zADataCollatorForTokenClassification.torch_call.<locals>.<listcomp>c �" g | ]
}� fd d�|� � D ��qS )c � i | ] \}}|� kr||�qS r r �r5 rD rR rg r r �128<dictcomp>K � zLDataCollatorForTokenClassification.torch_call.<locals>.<listcomp>.<dictcomp>�rN re rg r r r7 K � " r r` � input_ids� c s t | � j�r129| �� S t| �S r )rF rG �tolistr/ )Ztensor_or_iterable�rE r r �to_list\ s z>DataCollatorForTokenClassification.torch_call.<locals>.to_list�rightc s( g | ]}�|�� j g�t|� �qS r �rd �len�r5 r9 �r �sequence_lengthrt r r r7 b � �c s( g | ]}� j g�t|� �|� �qS r rv rx ry r r r7 f r{ r< )130rE r& r$ r] r^ r_ �shape� padding_siderL rV )r r r> Zno_labels_featuresrQ r} r )rh r rz rt rE r r E s4 "� �131�z-DataCollatorForTokenClassification.torch_callc s� dd l �d|d v r dnd� � |d v r � fdd�|D �nd }t�j|�j�j�j|d u r/dnd d�}|d u r9|S ��|d �jd ��jj}|d132krX��fd d�|D �|d< n ��fd d�|D �|d< �fd
d�|� � D �}|S )Nr r9 r> c rB r r re rg r r r7 q r8 z>DataCollatorForTokenClassification.tf_call.<locals>.<listcomp>r r` rp rq ru c �( g | ]}t |�� jg�t|� �qS r �r/ rd rw rx �r rz r r r7 � r{ c �( g | ]}� j g�t|� t|� �qS r �rd rw r/ rx r� r r r7 � r{ c s" i | ]
\}}|� j |� jd ��qS �r< �rZ rV rk �r r r rl � ro z>DataCollatorForTokenClassification.tf_call.<locals>.<dictcomp>)133rT r& r$ r] r^ r_ rZ r| r} rN �r r r> rQ r} r )rh r rz r r r m s0 "�134 � 135�z*DataCollatorForTokenClassification.tf_callc s� d|d v rdnd� � |d v r� fdd�|D �nd }t �j|�j�j�j|d u r+dnd d�}|d u r5|S t�|d �jd ��jj}|d136krT��fd d�|D �|d< n ��fd d�|D �|d< d
d� |� � D �}|S )Nr9 r r> c rB r r re rg r r r7 � r8 zADataCollatorForTokenClassification.numpy_call.<locals>.<listcomp>r r` rp rq ru c r~ r r rx r� r r r7 � r{ c r� r r� rx r� r r r7 � r{ c S s" i | ]
\}}|t j|t jd ��qS r� �r r[ rV rk r r r rl � ro zADataCollatorForTokenClassification.numpy_call.<locals>.<dictcomp>)137r& r$ r] r^ r_ r r[ r| r} rN r� r )rh r rz r r � s. "�138 � 139�z-DataCollatorForTokenClassification.numpy_call)r r r r- r140 r. r] r ra r r r^ r rI r_ rd r
r r r r r r r rb s 141 ( rb r_ c sD ddl �t| d tttjf�r�fdd�| D �} | d �d�� t� fdd�| D ��}|rC|du s6� | dkrCt| �j�sC�j | dd�S |j142du rRt d|j j
� d ���td143d� | D ��}|durm|| dkrm|| d | }| d �t| �|g|j�}t| �D ]!\}}|jd kr�|||d|jd �f< q~||||jd d�f< q~|S )
�_Collate `examples` into a batch, using the information in `tokenizer` for padding if necessary.r Nc � g | ]144}� j |� jd ��qS r� )rL rJ �r5 �ers r r r7 � � z(_torch_collate_batch.<locals>.<listcomp>c 3 s � | ]145}|� d �� kV qdS �r N��size�r5 �x��length_of_firstr r � <genexpr>� s � z'_torch_collate_batch.<locals>.<genexpr>)�dim�CYou are attempting to pad samples but the tokenizer you are using (� ) does not have a pad token.c s s � | ]}|� d �V qdS r� r� r� r r r r� � s � rq ru )rE rF r/ rY r rO r� �allrG rM � pad_tokenr � __class__r �maxZnew_fullrw � pad_token_id� enumerater} r| ��examplesr$ r_ �are_tensors_same_lengthr^ �result�i�exampler )r� rE r �_torch_collate_batch� s. 146��147r� c sL dd l � t| d ttf�r�fdd�| D �} t| d �� t� fdd�| D ��}|r;|d u s4� | dkr;�j| dd�S |jd u rJtd|j j148� d���t d d� | D ��}|d ure|| dkre|| d149 | }g }�� | d �}t
j|d ft
jd �}| D ]$}|jd
kr�|t|� |d< n|t|� |d< |��j|||jd�� qz�j|dd�S )Nr c r� r� r� r� r� r r r7 � r� z%_tf_collate_batch.<locals>.<listcomp>c 3 � � | ] }t |�� kV qd S r �rw r� r� r r r� � � � z$_tf_collate_batch.<locals>.<genexpr>�Zaxisr� r� c s � � | ]}t |�V qd S r r� r� r r r r� � � � rq r r< ru )r rq )r r )Zconstant_values)rT rF r/ rY rw r� rM r� r r� r r� �rankr �zeros�int32r} �appendr" r� ) r� r$ r_ r� r^ r� r� Zpaddingsr� r )r� r r �_tf_collate_batch� s4 150��151r� c s, t | d ttf�rdd� | D �} t| d �� t� fdd�| D ��}|r4|du s-� | dkr4tj| dd�S |jdu rCtd|j j152� d ���t d153d� | D ��}|dur^|| dkr^|| d | }tj t| �|f|j
| d jd �}t| �D ]!\}}|jd
kr�|||d|jd �f< qr||||jd d�f< qr|S )r� r c S s g | ]154}t j|t jd ��qS r� r� r� r r r r7 � r� z(_numpy_collate_batch.<locals>.<listcomp>c 3 r� r r� r� r� r r r� � r� z'_numpy_collate_batch.<locals>.<genexpr>Nr� r� r� c s r� r r� r� r r r r� r� rq )r| Z155fill_valuer= ru )rF r/ rY rw r� r rM r� r r� r r� �fullr� r= r� r} r| r� r r� r �_numpy_collate_batch� s* 156�� 157r� c @ s~ e Zd ZU dZeed< dZeee e158f ed< dZ e e
ed< dZe e
ed< dZe ed < d159eee ef fd d �Zd
d� ZdS )� DataCollatorForMultipleChoicea 160 Data collator that dynamically pads a batch of nested examples for multiple choice, so that all choices161 of all examples have the same length.162 163 Args:164 tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):165 The tokenizer used for encoding the data.166 padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):167 Select a strategy to pad the returned sequences according to the model's padding side and padding index168 among:169 170 - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence171 is provided).172 - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum173 acceptable input length for the model if that argument is not provided.174 - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different175 lengths).176 max_length (`int`, *optional*):177 Maximum length of the returned list and optionally padding length (see above).178 pad_to_multiple_of (`int`, *optional*):179 Pad the sequence to a multiple of the provided value.180 181 This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=182 7.5 (Volta).183 return_tensors (`str`, *optional*, defaults to `"pt"`):184 The type of Tensor to return. Allowable values are "np", "pt" and "tf".185 r$ Tr] Nr^ r_ r r
r� c s� dd l }d|d v r dnd��fdd�|D �}t|�� t|d d ��t�fdd�|D �g d �}| jj|| j| j| jd186d �}� �fd d
�|�� D �}|j ||j187d�|d< |S )Nr r9 r> c � g | ]}|� � ��qS r ��pop�r5 r� rg r r r7 @ � z<DataCollatorForMultipleChoice.torch_call.<locals>.<listcomp>rp c 3 s& � | ]� � fd d�t ��D �V qdS )c �" g | ]
� � fd d��� � D ��qS )c � i | ] \}}||� �qS r r rk �r� r r rl H � zQDataCollatorForMultipleChoice.torch_call.<locals>.<genexpr>.<listcomp>.<dictcomp>rn �r5 �r� r� r r7 H ro zFDataCollatorForMultipleChoice.torch_call.<locals>.<genexpr>.<listcomp>N��ranger� �� num_choicesr� r r� H s �$ z;DataCollatorForMultipleChoice.torch_call.<locals>.<genexpr>)�startr r` c s i | ] \}}||� � �d ��qS ������)�viewrk )�188batch_sizer� r r rl U s z<DataCollatorForMultipleChoice.torch_call.<locals>.<dictcomp>r< ) rE rw �sumr$ r" r] r^ r_ rN rL rV )r r� rE r> Z
flat_examplesrQ r )r� rh r� r r ; s$ �� z(DataCollatorForMultipleChoice.torch_callc s� dd l �d|d v r dnd��fdd�|D �}t|�� t|d d ���fdd�|D �}t|g �}| jj|| j| j| jdd �}� ��fd189d �|�� D �}�j |�j190d �|d< |S )
Nr r9 r> c r� r r� re rg r r r7 ] r� z9DataCollatorForMultipleChoice.tf_call.<locals>.<listcomp>rp c s" g | ]
� � fd d�t ��D ��qS )c r� )c r� r r rk r� r r rl a r� zODataCollatorForMultipleChoice.tf_call.<locals>.<listcomp>.<listcomp>.<dictcomp>rn r� �rf r� r r7 a ro zDDataCollatorForMultipleChoice.tf_call.<locals>.<listcomp>.<listcomp>r� r� r� r� r r7 ` s �r r` c s$ i | ]\}}|�� |� �d f��qS r� )�reshaperk )r� r� r r r rl m s $ z9DataCollatorForMultipleChoice.tf_call.<locals>.<dictcomp>r< ) rT rw r� r$ r" r] r^ r_ rN rZ rV )r r r> Zflattened_featuresrQ r )r� rh r� r r r Y s& 191�192�z%DataCollatorForMultipleChoice.tf_call)r r r r- r193 r. r] r ra r r r^ r rI r_ r
r/ r0 r r r r r r r r� s 194 r� c @ s� e Zd ZU dZeed< dZee ed< dZ e195e e e
f ed< dZee ed< dZee ed< d Zeed196< d Ze ed < dd
d�ZdS )�DataCollatorForSeq2Seqa� 197 Data collator that will dynamically pad the inputs received, as well as the labels.198 199 Args:200 tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):201 The tokenizer used for encoding the data.202 model ([`PreTrainedModel`], *optional*):203 The model that is being trained. If set and has the *prepare_decoder_input_ids_from_labels*, use it to204 prepare the *decoder_input_ids*205 206 This is useful when using *label_smoothing* to avoid calculating loss twice.207 padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):208 Select a strategy to pad the returned sequences (according to the model's padding side and padding index)209 among:210 211 - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single212 sequence is provided).213 - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum214 acceptable input length for the model if that argument is not provided.215 - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).216 max_length (`int`, *optional*):217 Maximum length of the returned list and optionally padding length (see above).218 pad_to_multiple_of (`int`, *optional*):219 If set will pad the sequence to a multiple of the provided value.220 221 This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=222 7.0 (Volta).223 label_pad_token_id (`int`, *optional*, defaults to -100):224 The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).225 return_tensors (`str`, *optional*, defaults to `"pt"`):226 The type of Tensor to return. Allowable values are "np", "pt" and "tf".227 r$ N�modelTr] r^ r_ rc rd r r
c sV |d u r�j }d|d v rdnd� � |d v r � fdd�|D �nd }|d ur1tdd� |D ��r1d }� fdd�|D �}t�j|�j�j�j|d �}�jd228u pR�jtjk}|d ur�|rst |d � t229�rit230|�|d< n]d d� |D �|d< nS�jtj ko}�jd u}|s�t d d� |D ��n�j��jd ur���j d
�j �j ��jj
�t |d � t231�r����fdd�|D �|d< n
���fdd�|D �|d< |�dd �d u�r|dkr�dd l}|j|d |jd�|d< n&|dkr�dd l} | j|d | jd�|d< ntj|d tjd�|d< nd |d< |d u�r)�jd u�r)t�jd��r)�jj|d d�}232|233|d< |S )Nr9 r r> c rB r r re rg r r r7 � r8 z3DataCollatorForSeq2Seq.__call__.<locals>.<listcomp>c s s � | ]}|d u V qd S r r rx r r r r� � r� z2DataCollatorForSeq2Seq.__call__.<locals>.<genexpr>c ri )c rj r r rk rg r r rl � rm z>DataCollatorForSeq2Seq.__call__.<locals>.<listcomp>.<dictcomp>rn re rg r r r7 � ro r` Fc S s g | ] }t �|g g��qS r )r � concatenaterx r r r r7 � r� c s r� r r� )r5 �lr r r r� � r� rq c sD g | ] }�d kr|�j g� t|� n �j g� t|� | �qS )ru rv rx ��max_label_lengthr} r r r r7 � s ���c 234 sh g | ]0}�d kr t �|t j�jg� t|� t jd�g�nt �t j�jg� t|� t jd�|g��qS )ru r< )r r� r[ rd rw rV rx r� r r r7 � s �� �� ���r r< r �%prepare_decoder_input_ids_from_labels)r> �decoder_input_ids)r
r� r& r$ r] r^ r_ r Z235DO_NOT_PADrF r/ Z236MAX_LENGTHr� r} r# rE rL rV rT Zconstantr r[ r� r! r� ) r r r
r> Znon_labels_featuresrQ Z237no_paddingZ max_paddingrE r r� r )rh r� r} r r r � sl "�238 239 ��� �240�241 242�z DataCollatorForSeq2Seq.__call__r )r r r r- r243 r. r� r r r] r ra r r r^ rI r_ rd r
r r r r r r� r s 244 ! r� c @ sX e Zd ZU dZeed< dZeed< dZeed< dZ e245e ed< d Z e ed246< d Z
e ed < d
Ze247e ed< dZeed< dZeed< d
Ze248e ed< dd� Zdd� Zdd� Zed3dd��Z
d3ded e249e d eeef fd d �Zd e e e e ee eef f d e eef fd!d"�Z d e e e e ee eef f d e eef fd#d$�Z
d4ded e250e d%e251e d eeef fd&d'�Z!d e e e e ee eef f d e eef fd(d)�Z"
d4ded e252e d%e253e d eeef fd*d+�Z#ed,e$j%e$j%eeef d e$j%e$j%e d ee$j%e$j%e e$j%e$j%e f fd-d.��Z&ed/e$j%e$j%e d0ed efd1d2��Z'd
S )5� DataCollatorForLanguageModelinga 254 Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they255 are not all of the same length.256 257 Args:258 tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):259 The tokenizer used for encoding the data.260 mlm (`bool`, *optional*, defaults to `True`):261 Whether or not to use masked language modeling. If set to `False`, the labels are the same as the inputs262 with the padding tokens ignored (by setting them to -100). Otherwise, the labels are -100 for non-masked263 tokens and the value to predict for the masked token.264 whole_word_mask (`bool`, *optional*, defaults to `False`):265 Whether or not to mask whole words instead of individual tokens.266 mlm_probability (`float`, *optional*, defaults to 0.15):267 The probability with which to (randomly) mask tokens in the input, when `mlm` is set to `True`.268 mask_replace_prob (`float`, *optional*, defaults to 0.8):269 The probability with which masked tokens are replaced by the tokenizer's mask token (e.g., `[MASK]`).270 Defaults to 0.8, meaning 80% of the masked tokens will be replaced with `[MASK]`.271 Only works when `mlm` is set to `True`.272 random_replace_prob (`float`, *optional*, defaults to 0.1):273 The probability with which masked tokens are replaced by random tokens from the tokenizer's vocabulary.274 Defaults to 0.1, meaning 10% of the masked tokens will be replaced with random tokens. The remaining275 masked tokens (1 - mask_replace_prob - random_replace_prob) are left unchanged.276 Only works when `mlm` is set to `True`.277 pad_to_multiple_of (`int`, *optional*):278 If set, will pad the sequence to a multiple of the provided value.279 return_tensors (`str`):280 The type of Tensor to return. Allowable values are "np", "pt" and "tf".281 seed (`int`, *optional*):282 The seed to use for the random number generator for masking. If not provided, the global RNG will be used.283 284 <Tip>285 286 For best performance, this data collator should be used with a dataset having items that are dictionaries or287 BatchEncoding, with the `"special_tokens_mask"` key, as returned by a [`PreTrainedTokenizer`] or a288 [`PreTrainedTokenizerFast`] with the argument `return_special_tokens_mask=True`.289 290 <Example Options and Expectations>291 292 1. Default Behavior:293 - `mask_replace_prob=0.8`, `random_replace_prob=0.1`.294 - Expect 80% of masked tokens replaced with `[MASK]`, 10% replaced with random tokens, and 10% left unchanged.295 296 2. All masked tokens replaced by `[MASK]`:297 - `mask_replace_prob=1.0`, `random_replace_prob=0.0`.298 - Expect all masked tokens to be replaced with `[MASK]`. No tokens are left unchanged or replaced with random tokens.299 300 3. No `[MASK]` replacement, only random tokens:301 - `mask_replace_prob=0.0`, `random_replace_prob=1.0`.302 - Expect all masked tokens to be replaced with random tokens. No `[MASK]` replacements or unchanged tokens.303 304 4. Balanced replacement:305 - `mask_replace_prob=0.5`, `random_replace_prob=0.4`.306 - Expect 50% of masked tokens replaced with `[MASK]`, 40% replaced with random tokens, and 10% left unchanged.307 308 Note:309 The sum of `mask_replace_prob` and `random_replace_prob` must not exceed 1. If their sum is less than 1, the310 remaining proportion will consist of masked tokens left unchanged.311 312 </Tip>313 r$ T�mlmF�whole_word_maskg333333�?�mlm_probability皙�����?�mask_replace_probg�������?�random_replace_probNr_ �tf_experimental_compiler r
�seedc C s* | j r'| jjd u r
td��| jd u s | jdk s | jdkr td��t| j�| _n| jr.td��| j| j dkr:td��| jdk sD| jdkrHtd��| jdk sR| jdkrVtd��t| j�| _t| j�| _| j rrdd l314}|j | j d d315�| _ | jr�| jj
st�d t� | jdk r�t�d d
� d| _d| _d | _d S )Nz�This tokenizer does not have a mask token which is necessary for masked language modeling. You should pass `mlm=False` to train on causal language modeling instead.r rq z*mlm_probability should be between 0 and 1.zmWhole word masking can only be used with mlm=True.If you want to use whole word masking, please set mlm=True.zHThe sum of mask_replace_prob and random_replace_prob should not exceed 1z,mask_replace_prob should be between 0 and 1.z.random_replace_prob should be between 0 and 1.T)Z jit_compilezcWhole word masking depends on offset mapping which is only natively available with fast tokenizers.zBRandom token replacement is not supported with whole word masking.z Setting mask_replace_prob to 1.)r� r$ �316mask_tokenr r� rK r� r� r� r� rT �function�tf_mask_tokensZis_fast�warnings�warn� UserWarning� generator)r r r r r �
__post_init__E sJ � � �317�318z-DataCollatorForLanguageModeling.__post_init__c C sL | j dkrdd l}|�� �|�S | j dkr dd l}|jj�|�S tj�|�S )Nr r r ) r
rE � GeneratorZ manual_seedrT �randomZ from_seedr Z default_rng)r r� rE r r r r �
get_generators s 319320 z-DataCollatorForLanguageModeling.get_generatorc C s` t �� jdkr| �| j�| _d S dd l}|jj� � }|d u r$d}t321|��| �| j|j �| _d S )N� MainProcessr )zZWorker process information is not available for seeding the generator. This may be becausezZyou are using multiprocessing without using a PyTorch DataLoader. The `seed` parameter canzVonly be used when using multiprocessing with a PyTorch DataLoader. Please either use azAsingle process or use a PyTorch DataLoader with multiple workers.) �mp�current_process�namer� r� r� rE �utils�dataZget_worker_infor �id)r rE Z worker_infoZ error_stringr r r �322create_rng s z*DataCollatorForLanguageModeling.create_rngc C sZ dd l }|�| |�}|r |�||�| dd� dk|j�S |�||j�| dd� dk|j�S )Nr rq )rT �fill�cast�uniformra r� )r| Z probabilityr� r Z prob_matrixr r r � tf_bernoulli� s323 "z,DataCollatorForLanguageModeling.tf_bernoulli�inputs�special_tokens_maskr' c C s� ddl }|�||j�}|�|�}| �|| j| j�| @ }|�||d�}| �|| j| j�|@ } |�| ||�}| jdks?| j dkrC||fS d| j }324| j |325 } | �|| | j�|@ | @ } | jrh| jj326|||jd�}
n327|j j328|||jd�}
|�| |
|�}||fS )zw329 Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.330 r Nrc rq ��maxvalr= ) rT r� r= r| r� r� r� �wherer� r� r� r� )r r� �331vocab_size�
mask_token_idr� r � input_shape�masked_indicesr> �indices_replaced�remaining_prob�random_replace_prob_scaled�indices_random� random_wordsr r r r� � s, 332333334���z.DataCollatorForLanguageModeling.tf_mask_tokensr� c s8 dd l }� jr� jd u r� �� t|d t�r"t� j|d� jd�}n dt |� j� jd�i}|�335dd �}� j r{|d u rW� fdd�|d � � �
� D �}|�|j||jd �|j�}n|�||j�}� j|�|d |j�|� jjt� j�d336�\|d< |d < |S |d }� jjd ur�|�|� jjkd |�}n|�|�}||d < |S )
Nr r �r
r_ rp �r_ r� c � g | ]337}� j j|d d��qS �T)Zalready_has_special_tokens�r$ Zget_special_tokens_mask�r5 �val�r r r r7 � � ��z;DataCollatorForLanguageModeling.tf_call.<locals>.<listcomp>r< )r� r� r� r> rc )rT r� r� r� rF r r& r$ r_ r� r� r� �numpyrr r� rZ rV ra r� r� rw r� r� �identity)r r� r rQ r� r> r r r r � s< �� 338� �� 339z'DataCollatorForLanguageModeling.tf_callc C s� | j r | jd u r | �� t|d t�r t| j|d| jd�}n dt|| j| jd�i}|� dd �}|� dd �}| j340rJ| j |d ||d�\|d< |d < |S |d � � }| jj
d ur^d341||| jj
k< ||d < |S ) Nr r r rp r r� �offset_mapping�r� r r> rc )r� r� r� rF r r& r$ r_ r� r� r� �torch_mask_tokens�cloner� �r r� rQ r� r r> r r r r s( �� 342� � z*DataCollatorForLanguageModeling.torch_callr c s� ddl }|�� }|�|j� j�}|du r � fdd�|�� D �}� jr7� �t|�t|��\}}|j ||j343d�}nt ||j �rA|�344� n|j ||j345d�}|j
|dd� |j|� jd��346� } � jrf|�� �|| ��} d || < |j|�|j� j�� jd��347� | @ }348� j�� jj�||349< � jd350ks�� jdkr�||fS d351� j } � j| } |j|�|j| �� jd��352� | @ |353 @ }
|jt� j�|j|j� jd �}||
||
< ||fS ) �S354 Prepare masked tokens inputs/labels for masked language modeling.355 r Nc r r r r r r r r7 + � �zEDataCollatorForLanguageModeling.torch_mask_tokens.<locals>.<listcomp>r< � ��value�r� rc rq �r= r� )rE r r� r| r� rr r� � _calc_word_ids_and_prob_mask�to_numpyrL ra rF rG � masked_fill_� bernoullir� Z356BoolTensor�_whole_word_maskr� r$ �convert_tokens_to_idsr� r� r rw rJ )r r� r� r rE r> �probability_matrix�word_ids� no_mask_maskr� r� r� r� r� r� r r r r
sN 357� �358359��360 ��361362 ��� z1DataCollatorForLanguageModeling.torch_mask_tokensc C s� | j r | jd u r | �� t|d t�r t| j|d| jd�}n dt|| j| jd�i}|� dd �}|� dd �}| j363rJ| j |d ||d�\|d< |d < |S t �
|d �}| jjd ur_d364||| jjk< ||d < |S ) Nr r r rp r r� r r r> rc )r� r� r� rF r r& r$ r_ r� r� r� �numpy_mask_tokensr �copyr� r r r r r ^ s( �� 365�� z*DataCollatorForLanguageModeling.numpy_callc s t �|�}t �|j� j�}|du r � fdd�|�� D �}� jr,� �t|�t|��\}}nt |t j366�r7|� t �nt j
|t d�}d||< � jrS� jjd||jd�� t �}n
t jjd||jd�� t �}� jri� �||�}d|| < � jr�� jjd� j|jd�� t �|@ } nt jjd� j|jd�� t �|@ } � jj|| < � jdks�� jdkr�||fS d� j }367� j|368 } � jr jjd| |jd�� t �|@ | @ } � jjdt� j�t �| �t jd �}
n#t jjd| |jd�� t �|@ | @ } t jjdt� j�t �| �t jd �}
|
|| < ||fS )369r Nc r r r r r r r r7 � r zEDataCollatorForLanguageModeling.numpy_mask_tokens.<locals>.<listcomp>r< r rq r� rc ��low�highr� r= )r r! r� r| r� rr r� r r rF rO �astypera r[ r� �binomialr� r r� r$ r� r� �integersrw Z
count_nonzerorV r )r r� r� r r> r r r r� r� r� r� r� r� r r r r } sh 370 371� 372�373 � � 374 � � 375376��������z1DataCollatorForLanguageModeling.numpy_mask_tokens�offsetsc C s� | dd�dd�df }| dd�dd�df }t j|ddd�}d|dd�df< t j|ddd�}d|dd�df< |�t�}| ||k|dkB @ }t j|dd�}d||< | }||fS )z�377 Map tokens to word ids and create mask of tokens to not mask.378 Tokens that are part of the same word will have the same word id and we will only379 set a mask probability for the first token of each word.380 Nr rq r� r� )r Zrollr% ra Zcumsum) r( r� Z token_startsZ381token_endsZprev_token_endsZprev_token_specialZ is_new_wordr Z prob_maskr r r r � s 382383z<DataCollatorForLanguageModeling._calc_word_ids_and_prob_maskr �maskc C s� t |�}| dk}| dd�dd�df | dd�ddd�f k|dd�dd�df @ |dd�ddd�f @ }tj||dd�ddd�f @ dd�S )z>384 Mask whole words based on word ids and mask.385 r� Nr r� )r r �any)r r) Z valid_idsZ same_wordr r r r � s X$z0DataCollatorForLanguageModeling._whole_word_maskr �NN)(r r r r- r386 r. r� ra r� r� r rK r� r� r_ rI r� r
r r� r� r� r� � staticmethodr� r rY r� r/ r r0 r r r
r r r rO r r r r r r r� � sv 387 > . ���388389�4/4. ����390391�4?"����392393�R��" � (r� c s2 e Zd ZdZdeeee eee ef f dee ef fdd�Z394deeee eee ef f dee ef fdd�Z deeee eee ef f dee ef fdd �Z d395d � Z
dd
ee fdd�Zdededeeef fdd�Zdededeeef fdd�Zdededeeef fdd�Z� fdd�Z� ZS )� DataCollatorForWholeWordMaska# 396 Data collator used for language modeling that masks entire words.397 398 - collates batches of tensors, honoring their tokenizer's pad_token399 - preprocesses batches for masked language modeling400 401 <Tip>402 403 This collator relies on details of the implementation of subword tokenization by [`BertTokenizer`], specifically404 that subword tokens are prefixed with *##*. For tokenizers that do not adhere to this scheme, this collator will405 produce an output that is roughly equivalent to [`.DataCollatorForLanguageModeling`].406 407 </Tip>r� r' c C � | j r | jd u r | �� t|d t�rdd� |D �}n |}dd� |D �}t|| j| jd�}g }|D ]C}g }t|d �D ]
}| j� |�}|�408|� q;d|v rlt|d �} t |d �}409t |410�D ]} | | v rkd|| || < q]|�411| �
|�� q1t|| j| jd�} | �|| �\}
}|
|d �S )412Nr c S r: �rp r r� r r r r7 r8 z;DataCollatorForWholeWordMask.torch_call.<locals>.<listcomp>c S � g | ]}d |i�qS r/ r r� r r r r7 r8 r rp � chinese_ref�##�rp r> )r� r� r� rF r r� r$ r_ rr �_convert_id_to_tokenr� rw r� r r
�r r� rp � batch_input� mask_labelsr� �413ref_tokensr� �token�ref_pos�len_seqr� �414batch_maskr� r> r r r r
�0 �415z'DataCollatorForWholeWordMask.torch_callc C s( dd l }| jr| jd u r| �� t|d t�r dd� |D �}n |}dd� |D �}t|| j| jd�}g }|D ]C}g }t |d �D ]
}| j�416|�} |� | � q?d|v rpt |d �}417t |d �} t
| �D ]} | |418v rod|| || < qa|� | �|�� q5t|| j| jd�}
| �|�||j�|
�\}}||d �S )419Nr c S r: r/ r r� r r r r7 7 r8 z8DataCollatorForWholeWordMask.tf_call.<locals>.<listcomp>c S r0 r/ r r� r r r r7 : r8 r rp r1 r2 r3 )rT r� r� r� rF r r� r$ r_ rr r4 r� rw r� r r� r� rV )r r� r rp r6 r7 r� r8 r� r9 r: r; r� r< r� r> r r r r . s2 �420z$DataCollatorForWholeWordMask.tf_callc C r. )421Nr c S r: r/ r r� r r r r7 X r8 z;DataCollatorForWholeWordMask.numpy_call.<locals>.<listcomp>c S r0 r/ r r� r r r r7 [ r8 r rp r1 r2 r3 )r� r� r� rF r r� r$ r_ rr r4 r� rw r� r r r5 r r r r Q r= z'DataCollatorForWholeWordMask.numpy_callc s� | j d u r t�� � � S | jdkr(dd l}|jt� �| jd�}� fdd�|D �S | jdkrTdd l}| j� d�d }|jj422j |� t� ��|d��
� �� }� fd d�|D �S | jd423kra| j�� � � S d S ) Nr r r c � g | ]}� | �qS r r �r5 r� �� cand_indexesr r r7 } r8 z9DataCollatorForWholeWordMask._shuffle.<locals>.<listcomp>r r )r� c r> r r r? r@ r r r7 � r8 r )r� r� �shuffler
rE �randpermrw r� rT Z424make_seedsZ experimentalZstateless_shuffler� r rr )r rA rE �indicesr r� r r@ r �_shuffler s 425426427428$429 �z%DataCollatorForWholeWordMask._shuffle� � input_tokensc
s> ddl m}m} t| j||f�st�d� g }t|�D ]&\}}|dks'|dkr(qt|�dkr;|� d�r;|d �430|� q|�431|g� q| � |�}t |t
dttt|�| j ����}g } t� � |D ]$}432t| �|kri nt| �t|433� |krtq_|434D ] } � �| � | �435| � qvq_t� �t| �kr�td ��� fd436d �tt|��D �} | S ) zM437 Get 0/1 labels for masked tokens with whole word mask proxy438 r )�
BertTokenizer�BertTokenizerFastz�DataCollatorForWholeWordMask is only suitable for BertTokenizer-like tokenizers. Please refer to the documentation for more information.z[CLS]z[SEP]rq r2 r� z?Length of covered_indexes is not equal to length of masked_lms.c s g | ]439}|� v r440d nd�qS )rq r r r? �Zcovered_indexesr r r7 � r� zADataCollatorForWholeWordMask._whole_word_mask.<locals>.<listcomp>)Z transformersrH rI rF r$ r� r� r� rw �441startswithr� rE �minr� rI �roundr� �set�addr r� )
r rG Zmax_predictionsrH rI rA r� r9 Znum_to_predictZ442masked_lmsZ index_set�indexr7 r rJ r r � s: �443" 444 �z-DataCollatorForWholeWordMask._whole_word_maskr� r7 c sT ddl }� jjdu rtd��|�� }|}� fdd�|�� D �}|j|j||jd�dd� � jj dur@|�445� jj �}|j|dd� |�� }d || < |j |�
|j� j�� jd446��� |@ } � j�� jj�|| < � jd kso� jdkrs||fS d � j }447� j|448 } |j |�
|j| �� jd449��� |@ | @ } |jt� j�|j|j� jd �}
|
| || < ||fS )
��450 Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. Set451 'mask_labels' means we use whole word mask (wwm), we directly mask idxs according to it's ref.452 r N��This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer.c r r r r r r r r7 � r zBDataCollatorForWholeWordMask.torch_mask_tokens.<locals>.<listcomp>r< r r rc r rq r )rE r$ r� r r rr r rL ra r� �eqr� r r� r| r� r� r r� r rw rJ )r r� r7 rE r> r r� � padding_maskr� r� r� r� r� r� r r r r
� sD �453� 454 ��455456 ��� z.DataCollatorForWholeWordMask.torch_mask_tokensc sZ ddl }|�|�}� jjdu rtd��|�|�}|�||j�}� fdd�|D �}||j||jd� @ }� jjdurD|� jj k}|| @ }|�457||d�}� � |� j � j
�|@ } |�458| � jj|�}� j dksi� jdkrm||fS d� j }459� j|460 } � � || � j
�|@ | @ } � j
r�� j
j|t� j�|jd �}
n
|jj|t� j�|jd �}
|�461| |
|�}||fS )462rQ r NrR c r r r r r r r r7 � r z?DataCollatorForWholeWordMask.tf_mask_tokens.<locals>.<listcomp>r< rc rq r� )rT r| r$ r� r r463 r� ra r� r� r� r� r� r� r� r� r� rw rV r� )r r� r7 r r� r> r� r� rT r� r� r� r� r� r r r r� � sB 464 �465466� 467468469��� z+DataCollatorForWholeWordMask.tf_mask_tokensc s� � j jdu r470td��t�|�}|�t�}� fdd�|�� D �}d|tj|td�< � j j dur8|� j j471k}d||< d|| < � j rQ� j j d� j
|jd ��t�|@ }ntjj d� j
|jd ��t�|@ }� j �� j j�||< � j
dksu� jdkry||fS d� j
}� j| } � j r�� j j d| |jd ��t�|@ | @ }472� j jdt� j �|jtjd473�} n!tjj d| |jd ��t�|@ | @ }474tjjdt� j �|jtjd475�} | |476 ||477< ||fS ) rQ NrR c r r r r r r r r7 7 r zBDataCollatorForWholeWordMask.numpy_mask_tokens.<locals>.<listcomp>r r< rc rq r� r" )r$ r� r r r! r% ra rr r[ r� r� r� r& r� r| r� r r� r' rw rV r ) r r� r7 r> r� r� rT r� r� r� r� r� r r r r ( sR �478479480� 481 � �482483��� ��� z.DataCollatorForWholeWordMask.numpy_mask_tokensc s. t �dt� t� j|i |�� d| _d| _d S )Nz�DataCollatorForWholeWordMask is deprecated and will be removed in a future version, you can now use DataCollatorForLanguageModeling with whole_word_mask=True instead.T)r� r� �
FutureWarning�super�__init__r� r� �r �args�kwargs�r� r r rW i s �484z%DataCollatorForWholeWordMask.__init__)rF )r r r r- r/ r rI r r0 r r r r rE r rY r
r� r rW �
__classcell__r r r[ r r- � s 44!4#! * 8 <Ar- c C s( t | t�r| S t| d�r| �� } | �� S )Nr )rF r/ r! r rr �r� r r r rr t s485 486487rr c C s4 t | tj�r| S t| d�r| �� �� �� S t�| �S )N�detach)rF r rO r! r^ �cpur r[ r] r r r r | s488 489490r c @ sZ e Zd ZdZdd� Zdeeeef deeef fdd�Z dede491eeef fd d492�Z d S ) �DataCollatorForSOPz�493 Data collator used for sentence order prediction task.494 495 - collates batches of tensors, honoring their tokenizer's pad_token496 - preprocesses batches for both masked language modeling and sentence order prediction497 c O s t �dt� d S )Nz�DataCollatorForSOP is deprecated and will be removed in a future version, you can now use DataCollatorForLanguageModeling instead.)r� r� rU rX r r r rW � s �zDataCollatorForSOP.__init__r� r' c 498 C s� dd l }ddlm} dd� |D �}t|| j�}| �|�\}}}dd� |D �}||d| jjd�}dd� |D �}|�|�} ||||| d �S )499Nr )� pad_sequencec S r: r/ r r� r r r r7 � r8 z/DataCollatorForSOP.__call__.<locals>.<listcomp>c S r: )�token_type_idsr r� r r r r7 � r8 T)Z batch_firstZ
padding_valuec S r: )�sentence_order_labelr r� r r r r7 � r8 )rp r> �attention_maskrb rc )rE Ztorch.nn.utils.rnnra r� r$ � mask_tokensr� rM )500r r� rE ra rp r> rd rb Zsop_label_listrc r r r r � s 501�zDataCollatorForSOP.__call__r� c
sX ddl }� jjdu rtd��|�� }|�|j� j�}� fdd�|�� D �}|j |j502||j d�dd� � jj durF|�
� jj�}|j |dd� |�|�� � }| �� }� jj durf|�
� jj�} |j | d d� d503|| < |�|�|jd ��� � |@ }504� j�� jj�||505< |�|�|jd ��� � |@ |506 @ } |jt� j�|j|jd�} | | || < |||fS )
z�507 Prepare masked tokens inputs/labels/attention_mask for masked language modeling: 80% MASK, 10% random, 10%508 original. N-gram not applied yet.509 r NrR c r r r r r r r r7 � r z2DataCollatorForSOP.mask_tokens.<locals>.<listcomp>r< r r g �?rc r� g �?)rE r$ r� r r r� r| r� rr r rL ra r� rS r� r rK r r rw rJ )
r r� rE r> r r� rT r� rd Zattention_padding_maskr� r� r� r r r re � s4 �510� 511 512 " 513z DataCollatorForSOP.mask_tokensN) r r r r- rW r/ r0 r r r rY re r r r r r` � s514 & r` c @ s4 e Zd ZU dZeed< dZeed< dZe ed< dZ515e ed< d e e
e e eee ef f d516ee ef fd d �Zd e e
e e eee ef f d517ee ef fd
d�Zd e e
e e eee ef f d518ee ef fdd�Zded519eeeeef fdd�Zded520eeeeef fdd�Zded521eeeeef fdd�ZdS )�*DataCollatorForPermutationLanguageModelingz�522 Data collator used for permutation language modeling.523 524 - collates batches of tensors, honoring their tokenizer's pad_token525 - preprocesses batches for permutation language modeling with procedures specific to XLNet526 r$ gUUUUUU�?�plm_probability� �max_span_lengthr r
r� r' c C �H t |d t�rdd� |D �}t|| j�}| �|�\}}}}||||d�S )Nr c S r: r/ r r� r r r r7 � r8 zIDataCollatorForPermutationLanguageModeling.torch_call.<locals>.<listcomp>�rp � perm_mask�target_mappingr> )rF r r� r$ r
�r r� rQ r� rl rm r> r r r r � �527 z5DataCollatorForPermutationLanguageModeling.torch_callc C rj )Nr c S r: r/ r r� r r r r7 � r8 zFDataCollatorForPermutationLanguageModeling.tf_call.<locals>.<listcomp>rk )rF r r� r$ r� rn r r r r � ro z2DataCollatorForPermutationLanguageModeling.tf_callc C rj )Nr c S r: r/ r r� r r r r7 � r8 zIDataCollatorForPermutationLanguageModeling.numpy_call.<locals>.<listcomp>rk )rF r r� r$ r rn r r r r � ro z5DataCollatorForPermutationLanguageModeling.numpy_callr� c s� ddl }� jjdu rtd��|�d�d dkrtd��|�� }|j|jd|jd�}|j |�d�|�d�|�d�f|j528d�}t |�d��D ]J}d}|�d�}||k r�|� d� j
d d��� } t| � j �}529||� |530| d d��� } d||| | | �f< ||5317 }||k sQ|�|�d��||< qD|j� fd d532�|�� D �|jd�} |j| d d � � jjdur�|�� jj�}
|j|
d d � |
| B }� jj||< d
|| < |j |�d�|�d�|�d�f|j533d�}t |�d��D ]V}|�|�d��}|�d|�d�d f��dd�}||� |�d�d � }|� |�dd��}|�|| || @ d� |�|�d�df�|�d|�d�f�k|| @ ||< q�|� � |||� � fS )�g 534 The masked tokens to be predicted for a particular sequence are determined by the following algorithm:535 536 0. Start from the beginning of the sequence by setting `cur_len = 0` (number of tokens processed so far).537 1. Sample a `span_length` from the interval `[1, max_span_length]` (length of span of tokens to be masked)538 2. Reserve a context of length `context_length = span_length / plm_probability` to surround span to be539 masked540 3. Sample a starting point `start_index` from the interval `[cur_len, cur_len + context_length -541 span_length]` and mask tokens `start_index:start_index + span_length`542 4. Set `cur_len = cur_len + context_length`. If `cur_len < max_len` (i.e. there are tokens remaining in the543 sequence to be processed), repeat from Step 1.544 r N��This tokenizer does not have a mask token which is necessary for permutation language modeling. Please add a mask token if you want to use this tokenizer.rq r ��This collator requires that sequence lengths be even to create a leakage-free perm_mask. Please see relevant comments in source code for details.r< )rq c r r r r r r r r7 0 r� zPDataCollatorForPermutationLanguageModeling.torch_mask_tokens.<locals>.<listcomp>r r rc r� ) rE r$ r� r r� r r� r| ra r� rW r� r ri rH rI rg �eyerL rr r r� rS r� r� �aranger� � transposerC �flattenrJ )r r� rE r> r� rm r� �cur_len�max_len� span_length�context_length� start_indexr� rT �
non_func_maskrl �545perm_indexr r r r
� sZ
��(546 �
� 547 548( &549�z<DataCollatorForPermutationLanguageModeling.torch_mask_tokensc 550 s� ddl }� jjdu rtd��|�|�d d dkr td��|�|�}tj|j�� dt d�}|�|�}tj551|d |d |d ftj d�}t t
|��D ]D}d}|�|�d } || k r�td� jd �}552t|553� j �} |td| |554 d � } d||| | |555 �f< || 7 }|| k sXt�|d �||< qI|j|�|�|j d�}|�|�}|�� fdd �|�� �� D ��}
|j|
|j d�}
||
@ }� jjdur�|� jjk}|| @ }||
B }|�|� jj|�}|�||d556�}g }t t
|��D ]P}|� |d �}|�|� |d |d d f��}|j � |�}|� |�|�d �}|�|| || @ d |�}|� |� ||d df�|� |d|d f�k|| @ � q�|j |dd
�}|�||j!�|�||j �||�||j!�fS )rp r Nrq rq r rr r< c r r r r r r r r7 � r zMDataCollatorForPermutationLanguageModeling.tf_mask_tokens.<locals>.<listcomp>rc r� r� r� )"rT r$ r� r r| r557 r r� Zas_listra r� rW r� rw r ri rI rg rs r� rZ r rr r� r� r� r� ru r� r� rB r� rM rV )r r� r r> r� Z labels_shaperm r� rw rx ry rz r{ r� rT r| rl r} r r r r� ^ sj
��558559"�
560561562��563 564565566 &��*z9DataCollatorForPermutationLanguageModeling.tf_mask_tokensc sR � j jdu r567td��|jd d dkrtd��t�|�}tj|jdtd�}tj|jd |jd |jd ftj d�}t568|jd �D ]C}d}|jd }||k ryt d� j d �}t
|� j �} |t d| | d � }569d|||570|571| �f< || 7 }||k sMt�|jd �||< q@tj� fdd �|�� D �td�} d|| < � j jdur�|� j jk} d572|| < | | B }
� j j||< d || < tj|jd |jd |jd ftj d�}t573|jd �D ]G}t�|jd �}|�d |jd d f�j}tj�|� |j�� }d ||| |
| @ < |�|jd df�|�d|jd f�k|| @ ||< q�|�tj �|||�tj �fS )
rp Nrq rq r r rr r< c r r r r r r r r7 � r� zPDataCollatorForPermutationLanguageModeling.numpy_mask_tokens.<locals>.<listcomp>r rc r� ) r$ r� r r| r r! r� ra r� rW r� r ri rI rg rs r[ rr r� r� r� rt r� �Tr� rB rv r% rV )r r� r> r� rm r� rw rx ry rz r{ r� rT r| rl r} r r r r � sX
��574(575�
� 576 577( 578&579� z<DataCollatorForPermutationLanguageModeling.numpy_mask_tokensN)r r r r- r580 r. rg rK ri rI r
r r/ r r r0 r r r rY r
r� r r r r r rf � s 581 444 c"krf c s6 e Zd ZdZddddd�� fdd�582Zd d d583�Z� ZS ) �DataCollatorWithFlatteninga� 584 Data collator used for padding free approach. Does the following:585 586 - concatenates the entire mini batch into single long sequence of shape [1, total_tokens]587 - uses `separator_id` to separate sequences within the concatenated `labels`, default value is -100588 - no padding will be added, returns `input_ids`, `labels` and `position_ids` by default589 - optionally returns the kwargs contained in FlashAttentionKwargs590 - optionally returns seq_idx indicating which sequence each token belongs to591 592 <Tip warning={true}>593 594 Using `DataCollatorWithFlattening` will flatten the entire mini batch into single long sequence.595 Make sure your attention computation is able to handle it!596 597 </Tip>598 Trc F)�return_position_ids� separator_id�return_flash_attn_kwargs�return_seq_idxc sL t � j|i |�� || _|| _|| _|| _h d�| _h d�| _ddh| _d S )N> r> � position_idsrp > r> �seq_idxr� rp � max_length_q� max_length_k) rV rW r� r� r� r� � _int_64_keys�_batch_dim_keys� _py_int_keys)r r� r� r� r� rY rZ r[ r r rW > s 599600z#DataCollatorWithFlattening.__init__Nc s |d u r| j }|d u r| j}d|d v }g g d�}| jr#|�dg i� | jr-|�dg i� | jr5dg}d}t|�D ]o\� }|d } |d | 7 < |r]|d |g|d dd � 7 < n|d |g| dd � 7 < | jr}|d ttt | ���7 < | jr�|d � fdd �tt | ��D �7 < | jr�|�601|d602 t | � � t |t | ��}q9| jr�| |d < |d < | |d
< |d< |dkr�dd l }603|604j
} |605j} |606j}
n|dkr�tj} tj} tj}
ntd|�d���|�� D ]"\}}|| jv r�|g}|| jv�r607| ||| jv �r| n|
d�||< q�|S )Nr> r r3 r� r� rp rq c s g | ]}� �qS r r )r5 �_�r� r r r7 h s z7DataCollatorWithFlattening.__call__.<locals>.<listcomp>r� Z
cu_seq_lens_qZ
cu_seq_lens_kr� r� r r z;return_tensors must be one of ("pt", "np"), return_tensors=z not supportedr< )r
r� r� �updater� r� r� r/ r� rw r� r� rE rL rV r� r r[ r rN r� r� r� )r r r
r� Zis_labels_providedrQ Z cu_seq_lensr^ �samplerp rE Zdata_clsZdtype_64Zdtype_32rD rR r r� r r P s^ 608$ &�609 �z#DataCollatorWithFlattening.__call__r+ )r r r r- rW r r\ r r r[ r r + s �r )r r )0�multiprocessingr� r� r� �collections.abcr � dataclassesr r �typingr r r r r r Ztokenization_utils_baser610 r� r ZInputDataClassr/ r0 r Z DataCollatorr r& r+ r, r( r) r* r\ rb rI r� r� r� r� r� r� r- rr rO r r` rf r r r r r �<module> sl $ ) 3 %& Z x R T