o ϯi<@sdZddlZddlZddlZddlmZddlmZmZddl Z ddl Z ddl Z eddZ eddZd d Zd d Zd dZGdddeZeZddeeeefdede jfddZdS)zp CLIP tokenizer Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI. N) lru_cache)UnionListcCstjtjtjtdS)Nzbpe_simple_vocab_16e6.txt.gz)ospathjoindirnameabspath__file__r r T/home/ubuntu/.local/lib/python3.10/site-packages/laion_clap/clap_module/tokenizer.py default_bpesr cCstttdtddtttdtddtttdtdd}|dd}d }td D]}||vrI|||d ||d7}q3d d |D}tt||S) a9 Returns list of utf-8 byte and a corresponding list of unicode strings. 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 signficant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup tables between utf-8 bytes and unicode strings. And avoids mapping to whitespace/control characters the bpe code barfs on. !~¡¬®ÿNrcSsg|]}t|qSr )chr).0nr r r (z$bytes_to_unicode..)listrangeordappenddictzip)bscsrbr r r bytes_to_unicodesN   r$cCs6t}|d}|ddD] }|||f|}q |S)zReturn set of symbol pairs in a word. Word is represented as tuple of symbols (symbols being variable-length strings). rrN)setadd)wordpairs prev_charcharr r r get_pairs,s r+cCs"t|}tt|}|SN)ftfyfix_texthtmlunescapestriptextr r r basic_clean8s r4cCstdd|}|}|S)Nz\s+ )resubr1r2r r r whitespace_clean>sr8c@s:eZdZedfdefddZddZddZd d ZdS) SimpleTokenizerNbpe_pathcsHt_ddjD_t|dd}|dd}dd|D}t t }|d d|D}|D] }| d |q;|sMd d g}nd d g|}| |tt|tt|_d djD_tt|tt|_dd|D_d |}t|dtj_tj_fdd|D_dS)NcSi|]\}}||qSr r rkvr r r Gz,SimpleTokenizer.__init__..utf-8 ricSsg|]}t|qSr )tuplesplit)rmerger r r rJsz,SimpleTokenizer.__init__..cSsg|]}|dqS)r )rr>r r r rLr cSr;r r r<r r r r?Ur@cSsi|]}||qSr r rtr r r r?W|z:|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+cg|]}j|qSr encoderrJselfr r r\r@)r$ byte_encoderitems byte_decodergzipopenreaddecoderDrvaluesrrextendrr rlenrPdecoder bpe_rankscacher6compile IGNORECASEpat vocab_sizeall_special_ids)rRr:special_tokensmergesvocabrEspecialr rQr __init__Es*      zSimpleTokenizer.__init__c sj|jvr j|St|dd|ddf}t|}|s#|dS t|fddd}|jvr4nu|\}}g}d}|t|krz|||} |||| | }Wn|||dYn3|||kr|t|dkr||d|kr||||d 7}n ||||d7}|t|ksBt|}|}t|dkrnt|}q$d |}|j|<|S) NrFTcsj|tdS)Ninf)r^getfloat)pairrQr r hrLz%SimpleTokenizer.bpe..)keyrrr5) r_rCr+minr^r\indexr[rr) rRtokenr'r(bigramfirstsecondnew_wordijr rQr bpe^sH     ,     zSimpleTokenizer.bpecsng}tt|}tj|D]#}dfdd|dD}|fdd | dDq|S)NrGc3|]}j|VqdSr,)rS)rr#rQr r z)SimpleTokenizer.encode..rAc3r|r,rO)r bpe_tokenrQr r r}r~r5) r8r4lowerr6findallrbrencoder[r{rD)rRr3 bpe_tokensrtr rQr rs &zSimpleTokenizer.encodecsDdfdd|D}tfdd|Djddddd }|S) NrGcrNr )r])rrtrQr r rr@z*SimpleTokenizer.decode..crNr )rU)rcrQr r rr@rAreplace)errorsrFr5)r bytearrayrYr)rRtokensr3r rQr rYs(zSimpleTokenizer.decode) __name__ __module__ __qualname__r strrir{rrYr r r r r9Ds ) r9Mtextscontext_lengthreturncst|tr|g}tjdtjdfdd|D}tjt||tjd}t|D]\}}t||kr;|d|}t |||dt|f<q+|S)a Returns the tokenized representation of given input string(s) Parameters ---------- texts : Union[str, List[str]] An input string or a list of input strings to tokenize context_length : int The context length to use; all CLIP models use 77 as the context length Returns ------- A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length] rHrIcs"g|] }gt|gqSr ) _tokenizerr)rr3 eot_token sot_tokenr r rs"ztokenize..)dtypeN) isinstancerrrPtorchzerosr\long enumeratetensor)rr all_tokensresultryrr rr tokenizes     r)r)__doc__rVr/r functoolsrtypingrrr-regexr6rr r$r+r4r8objectr9rrint LongTensorrr r r r s&    Q*