o i @s(ddlmZddlmZmZmZddlmZddlm Z m Z m Z m Z m Z ddlmZmZddlmZmZmZmZe eded ed efd d Ze eded ed efd dZd efddZd efddZd efddZde deded eeeffddZded eeeeeegeffddZ dS))Path)CallableOptionalTuple)SentencePieceProcessor)ModelRaggedchaindeserialize_attrserialize_attr)build_camembert_adapterbuild_xlmr_adapter)Tok2PiecesBackpropT Tok2PiecesInTTok2PiecesModelTTok2PiecesOutTvaluenamereturncCs|SN) to_protobuf_rrmodelrq/home/ubuntu/.local/lib/python3.10/site-packages/spacy_curated_transformers/tokenization/sentencepiece_encoder.py!serialize_sentencepiece_processor srcCs t|Sr)r from_protobufrrrrdeserialize_my_custom_classs rcC"t}t|t}|d||S)aConstruct a SentencePiece piece encoder model that accepts a list of token sequences or documents and returns a corresponding list of piece identifiers with CamemBERT post-processing applied. This model must be separately initialized using an appropriate loader. encoder)build_sentencepiece_encoder_v1r r set_refr!rrrr(build_camembert_sentencepiece_encoder_v1  r%cCs$tdtdtid}|d||S)zConstruct a SentencePiece piece encoder model that accepts a list of token sequences or documents and returns a corresponding list of piece identifiers. This model must be separately initialized using an appropriate loader. sentencepiece_encodersentencepiece_processor)forwardattrsr!)rsentencepiece_encoder_forwardrr#)rrrrr"'s r"cCr )aConstruct a SentencePiece piece encoder model that accepts a list of token sequences or documents and returns a corresponding list of piece identifiers with XLM-RoBERTa post-processing applied. This model must be separately initialized using an appropriate loader. r!)r"r rr#r$rrr#build_xlmr_sentencepiece_encoder_v18r&r,rXis_trainc Cs|jd}g}|D]C}|g}dg}|D]}|jrg} n||j} || |t| q|||d|t |j ||j |q |ddfS)Nr(r cSsgSrr)dYrrrgsz/sentencepiece_encoder_forward..) r*bos_idis_space encode_as_idstextextendappendleneos_idrops asarray1i) rr-r.spppiecesdoc doc_pieceslenstoken piece_idsrrrr+Fs(        r+pathcsdfdd }|S)zConstruct a callback that initializes a SentencePiece piece encoder model. path (Path): Path to the serialized SentencePiece model. Ncstt|jd<|S)Nr()r from_filestrr*)rr-YrBrrloadus z3build_sentencepiece_encoder_loader_v1..load)NNr)rBrGrrFr%build_sentencepiece_encoder_loader_v1js rHN)!pathlibrtypingrrrcurated_tokenizersr thinc.apirrr r r sentencepiece_adaptersr rtypesrrrrregisterrDbytesrrr%r"r,boolr+rHrrrrsN      $