o ߥi@s8dZddlmZmZmZddlZddlZddlZddlZddl Z ddl Z ddl Z ddl Z ddl mZddlmZddlmZddlmZeZejZejZdd Zd d Zee jjjed ZGd ddejZ GdddejZ!GdddejZ"GdddejZ#GdddejZ$GdddejZ%GdddejZ&GdddejZ'GdddejZ(Gdd d ejZ)Gd!d"d"ejZ*Gd#d$d$ejZ+Gd%d&d&ejZ,Gd'd(d(ejZ-Gd)d*d*ejZ.Gd+d,d,ejZ/Gd-d.d.ejZ0Gd/d0d0ejZ1Gd1d2d2ejZ2Gd3d4d4e2Z3Gd5d6d6ejZ4dS)7zPyTorch BERT model.)absolute_importdivisionprint_functionN)nn)SpaceTCnConfig) ModelFile) get_loggercCs |ddt|tdS)zImplementation of the gelu activation function. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results): 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) g??g@)torcherfmathsqrtxr]/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/models/nlp/space_T_cn/backbone.pygelu(s rcCs|t|SN)r sigmoidrrrrswish0sr)rrelurcs&eZdZdfdd ZddZZS) BertLayerNorm-q=cs<tt|tt||_tt||_ ||_ dS)zWConstruct a layernorm module in the TF style (epsilon inside the square root). N) superr__init__r Parameterr onesweightzerosbiasvariance_epsilon)self hidden_sizeeps __class__rrr9s zBertLayerNorm.__init__cCsN|jddd}||djddd}||t||j}|j||jS)NT)keepdim)meanpowr r r rr)r!rusrrrforwardAszBertLayerNorm.forward)r__name__ __module__ __qualname__rr- __classcell__rrr$rr7srcsFeZdZdZfddZ              dddZZS)BertEmbeddingszLConstruct the embeddings from word, position and token_type embeddings. cstt|t|j|j|_t|j|j|_ t|j |j|_ td|j|_ td|j|_ t|jdd|_t|j|_dS)N rr#)rr3rr Embedding vocab_sizer"word_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddingsmatch_type_embeddingstype_embeddingsr LayerNormDropouthidden_dropout_probdropoutr!configr$rrrLszBertEmbeddings.__init__Nc Csr|d}tj|tj|jd}|d|}|dur t|}||}||}|dur|durt ||  }t j |td|  }t|D]0\}}|D]'\}}|||}|dkriqZtj|||d|ddfdd|||ddf<qZqR||}||}|||}|dur||}||7}|dur||}||7}||}||}|S)N)dtypedevicerrGdim)sizer arangelongrH unsqueeze expand_as zeros_liker9nparraycpunumpytolistobject enumerateitemsr)r;r=r>r?r@rC)r! input_ids header_idstoken_type_idsmatch_type_idsl_hs header_lentype_idx col_dict_listidsheader_flatten_tokensheader_flatten_indexheader_flatten_outputtoken_column_idtoken_column_maskcolumn_start_indexheaders_length seq_length position_idswords_embeddingsheader_embeddingsbicol_dictkivilengthr;r= embeddingsr>r?rrrr-\sN              zBertEmbeddings.forward)NNNNNNNNNNNNNNr/r0r1__doc__rr-r2rrr$rr3Hs$ r3cs.eZdZfddZddZdddZZS) BertSelfAttentioncstt||j|jdkrtd|j|jf|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_dS)NrLThe hidden size (%d) is not a multiple of the number of attention heads (%d))rrvrr"num_attention_heads ValueErrorintattention_head_size all_head_sizerLinearquerykeyvaluerAattention_probs_dropout_probrCrDr$rrrs  zBertSelfAttention.__init__cC6|dd|j|jf}|j|}|ddddSNr&rr(rFrLrxr{viewpermuter!r new_x_shaperrrtranspose_for_scores  z&BertSelfAttention.transpose_for_scoresNcCs||}||}||}||}||}||} t||dd} | t|j } | |} t j dd| } | | } t| | } | dddd} | dd|jf} | j| } | S)Nr&rJrr(rFr)r~rrrr matmul transposer r r{rSoftmaxrCr contiguousrLr|r)r! hidden_statesattention_maskschema_link_matrixmixed_query_layermixed_key_layermixed_value_layer query_layer key_layer value_layerattention_scoresattention_probs context_layernew_context_layer_shaperrrr-s,          zBertSelfAttention.forwardrr/r0r1rrr-r2rrr$rrvs rvcs0eZdZdZfddZddZddZZS)!BertSelfAttentionWithRelationsRATzd Adapted from https://github.com/microsoft/rat-sql/blob/master/ratsql/models/transformer.py cstt||j|jdkrtd|j|jf|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_t d|j|j|_t d|j|j|_dSNrrw)rrrr"rxryrzr{r|rr}r~rrrArrCr7relation_k_embrelation_v_embrDr$rrrs,    z*BertSelfAttentionWithRelationsRAT.__init__cCrrrrrrrrrz6BertSelfAttentionWithRelationsRAT.transpose_for_scorescCsL||}||}||}||}||}||} ||} ||} t| | dd} |dd} | dddd}t|| }| dddd}| |t |j }||}t jdd|}||}t|| }| dddd}t||}| dddd}||}| dddd}|dd|jf}|j|}|S) 7 relation is [batch, seq len, seq len] r&rrr(rFrrJN)r~rrrrrr rrrr r r{rrrCrrLr|r)r!rrrelationrrr relation_k relation_vrrrr relation_k_t query_layer_trelation_attention_scoresrelation_attention_scores_tmerged_attention_scoresrrattention_probs_tcontext_relationcontext_relation_tmerged_context_layerrrrrr-sp         z)BertSelfAttentionWithRelationsRAT.forward)r/r0r1rurrr-r2rrr$rrs  rc,eZdZfddZddZddZZS))BertSelfAttentionWithRelationsTableformercstt||j|jdkrtd|j|jf|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_t d|j|_dSr)rrrr"rxryrzr{r|rr}r~rrrArrCr7schema_link_embeddingsrDr$rrr;s" z2BertSelfAttentionWithRelationsTableformer.__init__cCrrrrrrrrNrz>BertSelfAttentionWithRelationsTableformer.transpose_for_scorescCs||}||}||}||}|dddd}||}||} ||} t|| dd} | t |j } | |} | |} t j dd| } || } t| | }|dddd}|dd|jf}|j|}|S) rrrrFr(r&rrJN)r~rrrrrr rrr r r{rrrCrrLr|r)r!rrrrrrrrrrrrrrrrrrr-TsF      z1BertSelfAttentionWithRelationsTableformer.forwardrrrr$rr9s rc$eZdZfddZddZZS)BertSelfOutputcsBtt|t|j|j|_t|jdd|_t |j |_ dSNrr6) rrrrr}r"denserr@rArBrCrDr$rrrzBertSelfOutput.__init__cC&||}||}|||}|SrrrCr@r!r input_tensorrrrr-  zBertSelfOutput.forwardr.rrr$rr rc(eZdZdfdd ZdddZZS) BertAttentionnonecsRtt||dkrt||_|dkrt||_|dkr"t||_t||_dS)Nrratadd) rrrrvr!rrroutputr!rEschema_link_moduler$rrrs   zBertAttention.__init__NcCs||||}|||}|Sr)r!r)r!rrr self_outputattention_outputrrrr-s  zBertAttention.forwardrrr.rrr$rrs rcr)BertIntermediatecsHtt|t|j|j|_t|j t rt |j |_ dS|j |_ dSr) rrrrr}r"intermediate_sizer isinstance hidden_actstrACT2FNintermediate_act_fnrDr$rrrs  zBertIntermediate.__init__cCs||}||}|Sr)rrr!rrrrr-s  zBertIntermediate.forwardr.rrr$rrrrcr) BertOutputcsBtt|t|j|j|_t|jdd|_ t |j |_ dSr) rrrrr}rr"rrr@rArBrCrDr$rrrrzBertOutput.__init__cCrrrrrrrr-rzBertOutput.forwardr.rrr$rrrrcr) BertLayerrcs4tt|t||d|_t||_t||_dSNr) rrrr attentionr intermediaterrrr$rrrs  zBertLayer.__init__NcCs(||||}||}|||}|Sr)rrr)r!rrrrintermediate_output layer_outputrrrr-s   zBertLayer.forwardrrr.rrr$rrsrcs(eZdZfddZ dddZZS)SqlBertEncodercs8tt|t|tfddt|D|_dS)Ncg|]}tqSrcopydeepcopy.0_layerrr z+SqlBertEncoder.__init__..)rrrrr ModuleListranger)r!layersrEr$rrrs  zSqlBertEncoder.__init__TcCs:g}|jD]}|||}|r||q|s|||Srrappend)r!rroutput_all_encoded_layersall_encoder_layers layer_modulerrrr-s    zSqlBertEncoder.forward)Tr.rrr$rrs  rcs.eZdZdfdd Z   d ddZZS) BertEncoderrcs>tt|t||dtfddt|jD|_dS)Nrcrrrrrrrrrz(BertEncoder.__init__..) rrrrrrrnum_hidden_layersrrr$rrrs   zBertEncoder.__init__NTcCs<g}|jD]}||||}|r||q|s|||Srr)r!rrall_schema_link_matrixall_schema_link_maskrrrrrrr-s   zBertEncoder.forwardr)NNTr.rrr$rrs  rcr) BertPoolercs.tt|t|j|j|_t|_dSr) rrrrr}r"rTanh activationrDr$rrrszBertPooler.__init__cCs(|dddf}||}||}|S)Nr)rr)r!rfirst_token_tensor pooled_outputrrrr- s  zBertPooler.forwardr.rrr$rr rcr)BertPredictionHeadTransformcsRtt|t|j|j|_t|jt rt |jn|j|_ t |jdd|_ dSr)rrrrr}r"rrrrrtransform_act_fnrr@rDr$rrrs  z$BertPredictionHeadTransform.__init__cCs"||}||}||}|Sr)rrr@rrrrr-s   z#BertPredictionHeadTransform.forwardr.rrr$rrs rcr)BertLMPredictionHeadcsZtt|t||_tj|d|ddd|_||j_ t t |d|_ dS)NrFrF)r)rrrr transformrr}rLdecoderrrr rrr!rEbert_model_embedding_weightsr$rrr$s  zBertLMPredictionHead.__init__cCs||}|||j}|Sr)rrrrrrrr-2s zBertLMPredictionHead.forwardr.rrr$rr"s rcr)BertOnlyMLMHeadcstt|t|||_dSr)rrrr predictionsrr$rrr:s zBertOnlyMLMHead.__init__cC||}|Sr)r)r!sequence_outputprediction_scoresrrrr-? zBertOnlyMLMHead.forwardr.rrr$rr8rrcr)BertOnlyNSPHeadcs"tt|t|jd|_dSNr()rrrrr}r"seq_relationshiprDr$rrrFszBertOnlyNSPHead.__init__cCrr)r)r!rseq_relationship_scorerrrr-JrzBertOnlyNSPHead.forwardr.rrr$rrDs rcr)BertPreTrainingHeadscs.tt|t|||_t|jd|_dSr) rr rrrrr}r"rrr$rrrQs zBertPreTrainingHeads.__init__cCs||}||}||fSr)rr)r!rrrrrrrr-Ws  zBertPreTrainingHeads.forwardr.rrr$rr Orr cs:eZdZdZfddZddZe  d ddZZS) PreTrainedBertModelz An abstract class to handle weights initialization and a simple interface for dowloading and loading pretrained models. cs:tt|t|tstd|jj|jj||_ dS)NzParameter config in `{}(config)` should be an instance of class `SpaceTCnConfig`. To create a model from a Google pretrained model use `model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`) rr rrrryformatr%r/rE)r!rEinputskwargsr$rrrbs   zPreTrainedBertModel.__init__cCs|t|tjtjfr|jjjd|jjdnt|t r'|j j |jj dt|tjr:|j dur<|j j dSdSdS)z! Initialize the weights. g)r)stdr N) rrr}r7rdatanormal_rEinitializer_rangerrzero_fill_)r!modulerrrinit_bert_weightsls  z%PreTrainedBertModel.init_bert_weightsNcsZ|}d}tj|r |}n*t}td||t |d } | |Wdn1s0wY|}tj |t } t | } td| || g|Ri|} durftj |t} t| g}g}D]$}d}d|vr||dd}d|vr|dd}|r||||qnt||D] \}}||<qgggtd ddur_dfd d | t| d rd nddtdkrtd| jjttdddtd| jjttdkr#td| jjttdddtd| jjt|r+t || S)a Instantiate a PreTrainedBertModel from a pre-trained model file or a pytorch state dict. Download and cache the pre-trained model file if needed. Params: pretrained_model_name: either: - a str with the name of a pre-trained model to load selected in the list of: . `bert-base-uncased` . `bert-large-uncased` . `bert-base-cased` . `bert-large-cased` . `bert-base-multilingual-uncased` . `bert-base-multilingual-cased` . `bert-base-chinese` - a path or url to a pretrained model archive containing: . `bert_config.json` a configuration file for the model . `pytorch_model.bin` a PyTorch dump of a BertForPreTraining instance cache_dir: an optional path to a folder in which the pre-trained models will be cached. state_dict: an optional state dictionary (collections.OrderedDict object) to use instead of Google pre-trained models *inputs, **kwargs: additional input for the specific Bert class (ex: num_labels for BertForSequenceClassification) Nz)extracting archive file {} to temp dir {}zr:gzzModel config {}gammarbetar _metadatac shdurin |ddi}|||d|jD]\}}|dur1|||dq dS)Nr&T.)get_load_from_state_dict_modulesrY)rprefixlocal_metadatanamechild error_msgsloadmetadata missing_keys state_dictunexpected_keysrrr$s  z1PreTrainedBertModel.from_pretrained..loadbertzbert.)rrz7Weights of {} not initialized from pretrained model: {}z **********zWARNING missing weightsz0Weights from pretrained model not used in {}: {}zWARNING unexpected weights)r)!ospathisdirtempfilemkdtemploggerinfor tarfileopen extractalljoin CONFIG_NAMErfrom_json_file WEIGHTS_NAMEr r$keysreplacerzippopgetattrrrhasattrlenr%r/printshutilrmtree)clspretrained_model_namer' cache_dirr r resolved_archive_filetempdirserialization_dirarchive config_filerEmodel weights_pathold_keysnew_keysrnew_keyold_keyrr"rfrom_pretrainedzs                z#PreTrainedBertModel.from_pretrained)NN) r/r0r1rurr classmethodrPr2rrr$rr ]s  r csReZdZdZd fdd Z                   d ddZZS) SpaceTCnModelaJ SpaceTCnModel model ("Bidirectional Embedding Representations from a Transformer pretrained on STAR-T-CN"). Params: config: a SpaceTCnConfig class instance with the configuration to build a new model Inputs: `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`. Outputs: Tuple of (encoded_layers, pooled_output) `encoded_layers`: controled by `output_all_encoded_layers` argument: - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size], - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding to the last attention block of shape [batch_size, sequence_length, hidden_size], `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a classifier pretrained on top of the hidden state associated to the first character of the input (`CLF`) to train on the Next-Sentence task (see BERT's paper). Example: >>> # Already been converted into WordPiece token ids >>> input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]]) >>> input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]]) >>> token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]]) >>> config = modeling.SpaceTCnConfig(vocab_size_or_config_json_file=32000, hidden_size=768, >>> num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072) >>> model = modeling.SpaceTCnModel(config=config) >>> all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask) rcsBtt||t||_t||d|_t||_| |j dSr) rrRrr3rsrencoderrpoolerapplyrrr$rrrs  zSpaceTCnModel.__init__NTcCs|dur t|}|durt|}|dd}d|d}|||||||| | | | | |||||}|j|||||d}|d}||}|sM|d}||fS)NrFr(r g)rrrr&)r ones_likerQrOrsrSrT)r!rZr[token_order_idsr\rr]r^r_type_idsrarbrcrdrerfrgrhrirrrextended_attention_maskembedding_outputencoded_layersrrrrrr-s0      zSpaceTCnModel.forwardr)NNNNNNNNNNNNNNNNNNTrtrrr$rrRs.+ rRcr)Seq2SQLc stt|||_||_||_||_| |_||_||_ ||_ ||_ | |_ t || |_t || |_t |||_t || |_t |||||_t || || |_t |d|_t ||d|_t ||d|_t || d|_dS)NrrF)rr\riShSlsdrrH n_agg_ops n_cond_ops n_action_opsmax_select_num max_where_numrr} w_sss_model w_sse_model s_ht_model wc_ht_modelselect_agg_model w_op_model conn_model action_model slen_model wlen_model) r!r]r^lSr`rbrarcrdrerHr$rrr^s2  zSeq2SQL.__init__cCs ||_dSr)rH)r!rHrrr set_device|s zSeq2SQL.set_devicec* s"t|}t|} t||}t||}t||}t|g} g} g} g} g}g}g}g}t|D]\| d| d| dfddtj D}| dj fddtj D}||||ddt|dD|dg||}||fd dtd |dDfd dt| t D7|qPt j | t jd j} t j | t jd j} t j | t jd j} t j | t jd j} t j |t jd j}t j |t jd j}t j |t jd j}t j |t jd j}t |}t |jgj}t |jgj}t |jgj}t |jgj}t |j jgj}t |j jgj}t ||djgj}t || jgj}t|D]}||d | |||ddf<||d | |||ddf<||d | |||ddf<||d | |||ddf<||d ||ddf||ddddf<||d ||ddf||ddddf<||d ||ddf||ddddf<||d ||ddf||ddddf<q|d j|d} |d j|j d}!|d j|j d}"|d j|jd}#|d jj |d j}$|d j|d j dd}%|d j|d j dd}&|d j|d j dd}' |d j|d j dd}(!|d jj |d j"})|#|(|)| |%|$|&|'f|!|"ffS)NrFr(rcsg|]}d|qS)rri)elemrrrz#Seq2SQL.forward..rrcsg|] }dj|qS))rers)rur!rrrscSsg|]}|qSrrrsrrrrscsg|]}|qSrrrs) column_indexibrrrrvrcsg|]}dqS)rrr)indexrrrsrIr&)#maxrRrSrTrUrVrXrrrerdr>r tensorrNtorHrr] index_selectrlreshapernrormrcrkrbrirrfrgrhrjra)*r! wemb_layerl_nr^ start_indexrxtokensrbmax_l_nmax_l_hs conn_index slen_index wlen_index action_indexwhere_op_indexselect_agg_indexheader_pos_index query_indexwoisaiqilist header_indexbSconn_embslen_embwlen_emb action_embwo_embsa_embqv_embht_embrts_ccos_slens_wlens_action wo_output wc_outputwv_sswv_se sc_output sa_outputr)rxruryrzr!rr-s   $   "       .2    zSeq2SQL.forward)r/r0r1rrqr-r2rrr$rr\\s r\)5ru __future__rrrrr r*r@r1r-rUrRr r.modelscope.models.nlp.space_T_cn.configurationrmodelscope.utils.constantrmodelscope.utils.loggerrr/ CONFIGURATIONr5TORCH_MODEL_BIN_FILEr7rr functionalrrModulerr3rvrrrrrrrrrrrrrrr r rRr\rrrrsT    L:kJ   s