U
    5AfC@                     @   s   d Z ddlZddlZddlmZ ddlmZmZmZ ddl	Z
ddlmZmZ ddlmZ eeZdd	d
Ze dd Zdd ZG dd deZdS )z!Tokenization classes for RoBERTa.    N)	lru_cache)ListOptionalTuple   )
AddedTokenPreTrainedTokenizer)loggingz
vocab.jsonz
merges.txt)
vocab_filemerges_filec                  C   s   t ttdtdd t ttdtdd  t ttdtdd  } | dd }d	}td
D ],}|| krf| | |d
|  |d7 }qfdd |D }tt| |S )a8  
    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
    characters the bpe code barfs on.

    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
    tables between utf-8 bytes and unicode strings.
    !~      ¡   ¬   ®   ÿNr      c                 S   s   g | ]}t |qS  )chr).0nr   r   T/tmp/pip-unpacked-wheel-zw5xktn0/transformers/models/roberta/tokenization_roberta.py
<listcomp>9   s     z$bytes_to_unicode.<locals>.<listcomp>)listrangeordappenddictzip)bscsr   br   r   r   bytes_to_unicode$   s    L

r#   c                 C   s6   t  }| d }| dd D ]}|||f |}q|S )z
    Return set of symbol pairs in a word.

    Word is represented as tuple of symbols (symbols being variable-length strings).
    r   r   N)setadd)wordpairsZ	prev_charcharr   r   r   	get_pairs=   s    r)   c                	       s  e Zd ZdZeZddgZd) fdd	Zedd Z	dd Z
dd Zdd Zdd Zdd Zdd Zd*eee ee dddZd+ee eee  ee dd d!Zd,ee eee  eee d" fd#d$Zd-ee eee  ee dd%d&Zd.d'd(Z  ZS )/RobertaTokenizera  
    Constructs a RoBERTa tokenizer, derived from the GPT-2 tokenizer, using byte-level Byte-Pair-Encoding.

    This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
    be encoded differently whether it is at the beginning of the sentence (without space) or not:

    ```python
    >>> from transformers import RobertaTokenizer

    >>> tokenizer = RobertaTokenizer.from_pretrained("FacebookAI/roberta-base")
    >>> tokenizer("Hello world")["input_ids"]
    [0, 31414, 232, 2]

    >>> tokenizer(" Hello world")["input_ids"]
    [0, 20920, 232, 2]
    ```

    You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you
    call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance.

    <Tip>

    When used with `is_split_into_words=True`, this tokenizer will add a space before each word (even the first one).

    </Tip>

    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
    this superclass for more information regarding those methods.

    Args:
        vocab_file (`str`):
            Path to the vocabulary file.
        merges_file (`str`):
            Path to the merges file.
        errors (`str`, *optional*, defaults to `"replace"`):
            Paradigm to follow when decoding bytes to UTF-8. See
            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
        bos_token (`str`, *optional*, defaults to `"<s>"`):
            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.

            <Tip>

            When building a sequence using special tokens, this is not the token that is used for the beginning of
            sequence. The token used is the `cls_token`.

            </Tip>

        eos_token (`str`, *optional*, defaults to `"</s>"`):
            The end of sequence token.

            <Tip>

            When building a sequence using special tokens, this is not the token that is used for the end of sequence.
            The token used is the `sep_token`.

            </Tip>

        sep_token (`str`, *optional*, defaults to `"</s>"`):
            The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
            sequence classification or for a text and a question for question answering. It is also used as the last
            token of a sequence built with special tokens.
        cls_token (`str`, *optional*, defaults to `"<s>"`):
            The classifier token which is used when doing sequence classification (classification of the whole sequence
            instead of per-token classification). It is the first token of the sequence when built with special tokens.
        unk_token (`str`, *optional*, defaults to `"<unk>"`):
            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
            token instead.
        pad_token (`str`, *optional*, defaults to `"<pad>"`):
            The token used for padding, for example when batching sequences of different lengths.
        mask_token (`str`, *optional*, defaults to `"<mask>"`):
            The token used for masking values. This is the token used when training this model with masked language
            modeling. This is the token which the model will try to predict.
        add_prefix_space (`bool`, *optional*, defaults to `False`):
            Whether or not to add an initial space to the input. This allows to treat the leading word just as any
            other word. (RoBERTa tokenizer detect beginning of words by the preceding space).
    Z	input_idsZattention_maskreplace<s></s><unk><pad><mask>Fc                    s  t |trt|dddn|}t |	tr4t|	dddn|	}	t |trPt|dddn|}t |trlt|dddn|}t |trt|dddn|}t |trt|dddn|}t |
trt|
ddddn|
}
t|dd}t|| _W 5 Q R X dd | j D | _|| _	t
 | _d	d | j D | _t|dd}| d
dd }W 5 Q R X dd |D }tt|tt|| _i | _|| _td| _t jf |||||||	|
|d	| d S )NF)lstriprstripT)r1   r2   
normalizedutf-8encodingc                 S   s   i | ]\}}||qS r   r   r   kvr   r   r   
<dictcomp>   s      z-RobertaTokenizer.__init__.<locals>.<dictcomp>c                 S   s   i | ]\}}||qS r   r   r7   r   r   r   r:      s      
r   c                 S   s   g | ]}t | qS r   )tuplesplit)r   merger   r   r   r      s     z-RobertaTokenizer.__init__.<locals>.<listcomp>zJ's|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+)	errors	bos_token	eos_token	unk_token	sep_token	cls_token	pad_token
mask_tokenadd_prefix_space)
isinstancestrr   openjsonloadencoderitemsdecoderr@   r#   byte_encoderbyte_decoderreadr>   r   r   r   len	bpe_rankscacherH   recompilepatsuper__init__)selfr
   r   r@   rA   rB   rD   rE   rC   rF   rG   rH   kwargsZvocab_handleZmerges_handleZ
bpe_merges	__class__r   r   r[      sH     
zRobertaTokenizer.__init__c                 C   s
   t | jS N)rT   rN   r\   r   r   r   
vocab_size   s    zRobertaTokenizer.vocab_sizec                 C   s   t | j }|| j |S r`   )r   rN   copyupdateZadded_tokens_encoder)r\   Zvocabr   r   r   	get_vocab   s    zRobertaTokenizer.get_vocabc           
         sd  | j kr j | S t|}t|}|s,|S t| fddd}| jkrNqL|\}}g }d}|t|k r"z|||}	W n, tk
r   |||d   Y q"Y nX ||||	  |	}|| |kr
|t|d k r
||d  |kr
|	||  |d7 }q^|	||  |d7 }q^t|}|}t|dkrBqLq,t|}q,d
|}| j |< |S )Nc                    s    j | tdS )Ninf)rU   getfloat)pairra   r   r   <lambda>       z&RobertaTokenizer.bpe.<locals>.<lambda>keyr   r       )rV   r=   r)   minrU   rT   index
ValueErrorextendr   join)
r\   tokenr&   r'   ZbigramfirstsecondZnew_wordijr   ra   r   bpe   sB    


2




zRobertaTokenizer.bpec                    sZ   g }t  j|D ]B}d fdd|dD }|dd  |dD  q|S )zTokenize a string. c                 3   s   | ]} j | V  qd S r`   )rQ   )r   r"   ra   r   r   	<genexpr>  s    z-RobertaTokenizer._tokenize.<locals>.<genexpr>r4   c                 s   s   | ]
}|V  qd S r`   r   )r   Z	bpe_tokenr   r   r   r|     s     ro   )rW   findallrY   rt   encoders   rz   r>   )r\   text
bpe_tokensru   r   ra   r   	_tokenize  s    "zRobertaTokenizer._tokenizec                 C   s   | j || j | jS )z0Converts a token (str) in an id using the vocab.)rN   rg   rC   )r\   ru   r   r   r   _convert_token_to_id  s    z%RobertaTokenizer._convert_token_to_idc                 C   s   | j |S )z=Converts an index (integer) in a token (str) using the vocab.)rP   rg   )r\   rq   r   r   r   _convert_id_to_token  s    z%RobertaTokenizer._convert_id_to_tokenc                    s0   d |}t fdd|D jd jd}|S )z:Converts a sequence of tokens (string) in a single string.r{   c                    s   g | ]} j | qS r   )rR   )r   cra   r   r   r      s     z=RobertaTokenizer.convert_tokens_to_string.<locals>.<listcomp>r4   )r@   )rt   	bytearraydecoder@   )r\   tokensr   r   ra   r   convert_tokens_to_string  s    
"z)RobertaTokenizer.convert_tokens_to_stringN)save_directoryfilename_prefixreturnc           
   	   C   s(  t j|s"td| d d S t j||r6|d ndtd  }t j||rX|d ndtd  }t|ddd	$}|t	j
| jd
dddd  W 5 Q R X d}t|ddd	j}|d t| j dd dD ]B\}}	||	krtd| d |	}|d|d  |d7 }qW 5 Q R X ||fS )NzVocabulary path (z) should be a directory-r{   r
   r   wr4   r5   rn   TF)indent	sort_keysensure_asciir;   r   z#version: 0.2
c                 S   s   | d S )Nr   r   )kvr   r   r   rj   4  rk   z2RobertaTokenizer.save_vocabulary.<locals>.<lambda>rl   zSaving vocabulary to zZ: BPE merge indices are not consecutive. Please check that the tokenizer is not corrupted!ro   r   )ospathisdirloggererrorrt   VOCAB_FILES_NAMESrK   writerL   dumpsrN   sortedrU   rO   warning)
r\   r   r   r
   Z
merge_filefrq   writerr   Ztoken_indexr   r   r   save_vocabulary#  s2      (

z RobertaTokenizer.save_vocabulary)token_ids_0token_ids_1r   c                 C   sD   |dkr| j g| | jg S | j g}| jg}|| | | | | S )a  
        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
        adding special tokens. A RoBERTa sequence has the following format:

        - single sequence: `<s> X </s>`
        - pair of sequences: `<s> A </s></s> B </s>`

        Args:
            token_ids_0 (`List[int]`):
                List of IDs to which the special tokens will be added.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.

        Returns:
            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
        N)cls_token_idsep_token_id)r\   r   r   clssepr   r   r    build_inputs_with_special_tokens@  s
    z1RobertaTokenizer.build_inputs_with_special_tokens)r   r   already_has_special_tokensr   c                    sh   |rt  j||ddS |dkr8dgdgt|  dg S dgdgt|  ddg dgt|  dg S )a  
        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
        special tokens using the tokenizer `prepare_for_model` method.

        Args:
            token_ids_0 (`List[int]`):
                List of IDs.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.
            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
                Whether or not the token list is already formatted with special tokens for the model.

        Returns:
            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
        T)r   r   r   Nr   r   )rZ   get_special_tokens_maskrT   )r\   r   r   r   r^   r   r   r   Y  s      z(RobertaTokenizer.get_special_tokens_maskc                 C   sP   | j g}| jg}|dkr.t|| | dg S t|| | | | | dg S )a  
        Create a mask from the two sequences passed to be used in a sequence-pair classification task. RoBERTa does not
        make use of token type ids, therefore a list of zeros is returned.

        Args:
            token_ids_0 (`List[int]`):
                List of IDs.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.

        Returns:
            `List[int]`: List of zeros.
        Nr   )r   r   rT   )r\   r   r   r   r   r   r   r   $create_token_type_ids_from_sequencest  s
    z5RobertaTokenizer.create_token_type_ids_from_sequencesc                 K   s>   | d| j}|s|r6t|dkr6|d  s6d| }||fS )NrH   r   ro   )poprH   rT   isspace)r\   r   Zis_split_into_wordsr]   rH   r   r   r   prepare_for_tokenization  s     z)RobertaTokenizer.prepare_for_tokenization)	r+   r,   r-   r-   r,   r.   r/   r0   F)N)N)NF)N)F)__name__
__module____qualname____doc__r   Zvocab_files_namesZmodel_input_namesr[   propertyrb   re   rz   r   r   r   r   rJ   r   r   r   r   intr   boolr   r   r   __classcell__r   r   r^   r   r*   K   sV   M         <
*
  
    
   
r*   )r   rL   r   	functoolsr   typingr   r   r   regexrW   Ztokenization_utilsr   r   utilsr	   Z
get_loggerr   r   r   r#   r)   r*   r   r   r   r   <module>   s   

