#! python3 # -*- encoding: utf-8 -*- from copy import deepcopy from torch.nn.init import xavier_uniform_ import torch.nn.functional as F from torch.nn import Parameter from torch.nn.init import normal_ import torch.utils.checkpoint from torch import Tensor, device from .TAAS_utils import * from transformers.modeling_utils import ModuleUtilsMixin from transformers import AutoTokenizer, AutoModel, BertTokenizer from .graphormer import Graphormer3D import pickle import torch import sys from .ner_model import NER_model import numpy as np from .htc_loss import HTCLoss from transformers.utils.hub import cached_file remap_code_2_chn_file_path = cached_file( 'Cainiao-AI/TAAS', 'remap_code_2_chn.pkl' ) s2_label_dict_remap = { 0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: 'a', 11: 'b', 12: 'c', 13: 'd', 14: 'e', 15: 'f'} class StellarEmbedding(nn.Module): """Construct the embeddings from word, position and token_type embeddings.""" def __init__(self, config): super().__init__() self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) self.ner_type_embeddings = nn.Embedding(10, config.hidden_size) self.use_task_id = config.use_task_id if config.use_task_id: self.task_type_embeddings = nn.Embedding(config.task_type_vocab_size, config.hidden_size) # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load # any TensorFlow checkpoint file self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.dropout = nn.Dropout(config.hidden_dropout_prob) # position_ids (1, len position emb) is contiguous in memory and exported when serialized self.position_embedding_type = getattr(config, "position_embedding_type", "absolute") self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))) self.register_buffer("token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False) self._reset_parameters() def forward( self, input_ids: Optional[torch.LongTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, ner_type_ids: Optional[torch.LongTensor] = None, task_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, past_key_values_length: int = 0, ) -> torch.Tensor: if input_ids is not None: input_shape = input_ids.size() else: input_shape = inputs_embeds.size()[:-1] seq_length = input_shape[1] if position_ids is None: position_ids = self.position_ids[:, past_key_values_length: seq_length + past_key_values_length] # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves # issue #5664 if token_type_ids is None: if hasattr(self, "token_type_ids"): buffered_token_type_ids = self.token_type_ids[:, :seq_length] buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length) token_type_ids = buffered_token_type_ids_expanded else: token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) token_type_embeddings = self.token_type_embeddings(token_type_ids) if ner_type_ids is not None: ner_type_embeddings = self.ner_type_embeddings(ner_type_ids) embeddings = inputs_embeds + token_type_embeddings + ner_type_embeddings else: embeddings = inputs_embeds + token_type_embeddings if self.position_embedding_type == "absolute": position_embeddings = self.position_embeddings(position_ids) embeddings += position_embeddings # add `task_type_id` for ERNIE model if self.use_task_id: if task_type_ids is None: task_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) task_type_embeddings = self.task_type_embeddings(task_type_ids) embeddings += task_type_embeddings embeddings = self.LayerNorm(embeddings) embeddings = self.dropout(embeddings) return embeddings def _reset_parameters(self): for p in self.parameters(): if p.dim() > 1: normal_(p, mean=0.0, std=0.02) def set_pretrained_weights(self, path): pre_train_weights = torch.load(path, map_location=torch.device('cpu')) new_weights = dict() for layer in self.state_dict().keys(): if layer == 'position_ids': new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_ids'] elif layer == 'word_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.word_embeddings.weight'] elif layer == 'position_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_embeddings.weight'] elif layer == 'token_type_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.token_type_embeddings.weight'] elif layer == 'task_type_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.task_type_embeddings.weight'] elif layer == 'LayerNorm.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.weight'] elif layer == 'LayerNorm.bias': new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.bias'] else: new_weights[layer] = self.state_dict()[layer] self.load_state_dict(new_weights) def save_weights(self, path): torch.save(self.state_dict(), path) def load_weights(self, path): self.load_state_dict(torch.load(path)) # Copied from transformers.models.bert.modeling_bert.BertLayer class StellarLayer(nn.Module): def __init__(self, config): super().__init__() self.chunk_size_feed_forward = config.chunk_size_feed_forward self.seq_len_dim = 1 self.attention = ErnieAttention(config) self.is_decoder = config.is_decoder self.add_cross_attention = config.add_cross_attention if self.add_cross_attention: if not self.is_decoder: raise ValueError(f"{self} should be used as a decoder model if cross attention is added") self.crossattention = ErnieAttention(config, position_embedding_type="absolute") self.intermediate = ErnieIntermediate(config) self.output = ErnieOutput(config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, ) -> Tuple[torch.Tensor]: # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None self_attention_outputs = self.attention( hidden_states, attention_mask, head_mask, output_attentions=output_attentions, past_key_value=self_attn_past_key_value, ) attention_output = self_attention_outputs[0] # if decoder, the last output is tuple of self-attn cache if self.is_decoder: outputs = self_attention_outputs[1:-1] present_key_value = self_attention_outputs[-1] else: outputs = self_attention_outputs[1:] # add self attentions if we output attention weights cross_attn_present_key_value = None if self.is_decoder and encoder_hidden_states is not None: if not hasattr(self, "crossattention"): raise ValueError( f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers" " by setting `config.add_cross_attention=True`" ) # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None cross_attention_outputs = self.crossattention( attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, cross_attn_past_key_value, output_attentions, ) attention_output = cross_attention_outputs[0] outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights # add cross-attn cache to positions 3,4 of present_key_value tuple cross_attn_present_key_value = cross_attention_outputs[-1] present_key_value = present_key_value + cross_attn_present_key_value layer_output = apply_chunking_to_forward( self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output ) outputs = (layer_output,) + outputs # if decoder, return the attn key/values as the last output if self.is_decoder: outputs = outputs + (present_key_value,) return outputs def feed_forward_chunk(self, attention_output): intermediate_output = self.intermediate(attention_output) layer_output = self.output(intermediate_output, attention_output) return layer_output class StellarEncoder(nn.Module): def __init__(self, config): super().__init__() self.config = config self.layer = nn.ModuleList([StellarLayer(config) for _ in range(config.num_hidden_layers)]) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = False, output_hidden_states: Optional[bool] = False, return_dict: Optional[bool] = True, ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None next_decoder_cache = () if use_cache else None for i, layer_module in enumerate(self.layer): if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) layer_head_mask = head_mask[i] if head_mask is not None else None past_key_value = past_key_values[i] if past_key_values is not None else None if self.gradient_checkpointing and self.training: if use_cache: logger.warning( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs, past_key_value, output_attentions) return custom_forward layer_outputs = torch.utils.checkpoint.checkpoint( create_custom_forward(layer_module), hidden_states, attention_mask, layer_head_mask, encoder_hidden_states, encoder_attention_mask, ) else: layer_outputs = layer_module( hidden_states, attention_mask, layer_head_mask, encoder_hidden_states, encoder_attention_mask, past_key_value, output_attentions, ) hidden_states = layer_outputs[0] if use_cache: next_decoder_cache += (layer_outputs[-1],) if output_attentions: all_self_attentions = all_self_attentions + (layer_outputs[1],) if self.config.add_cross_attention: all_cross_attentions = all_cross_attentions + (layer_outputs[2],) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) if not return_dict: return tuple( v for v in [ hidden_states, next_decoder_cache, all_hidden_states, all_self_attentions, all_cross_attentions, ] if v is not None ) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=next_decoder_cache, hidden_states=all_hidden_states, attentions=all_self_attentions, cross_attentions=all_cross_attentions, ) # Copied from transformers.models.bert.modeling_bert.BertPooler class StellarPooler(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.activation = nn.Tanh() def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # We "pool" the model by simply taking the hidden state corresponding # to the first token. first_token_tensor = hidden_states[:, 0] pooled_output = self.dense(first_token_tensor) pooled_output = self.activation(pooled_output) return pooled_output class StellarModel(nn.Module): """ """ def __init__(self, config, add_pooling_layer=True): super().__init__() self.config = config self.encoder = StellarEncoder(config) self.pooler = StellarPooler(config) if add_pooling_layer else None # Initialize weights and apply final processing self._reset_parameters() # Copied from transformers.models.bert.modeling_bert.BertModel._prune_heads def _prune_heads(self, heads_to_prune): """ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel """ for layer, heads in heads_to_prune.items(): self.encoder.layer[layer].attention.prune_heads(heads) def forward( self, h_input, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, task_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: r""" encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is configured as a decoder. encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if self.config.is_decoder: use_cache = use_cache if use_cache is not None else self.config.use_cache else: use_cache = False if input_ids is not None and inputs_embeds is not None: raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") elif input_ids is not None: input_shape = input_ids.size() elif inputs_embeds is not None: input_shape = inputs_embeds.size()[:-1] else: raise ValueError("You have to specify either input_ids or inputs_embeds") batch_size, seq_length = input_shape device = input_ids.device if input_ids is not None else inputs_embeds.device # past_key_values_length past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 if attention_mask is None: attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) if token_type_ids is None: if hasattr(self.embeddings, "token_type_ids"): buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) token_type_ids = buffered_token_type_ids_expanded else: token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] # ourselves in which case we just need to make it broadcastable to all heads. extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) # If a 2D or 3D attention mask is provided for the cross-attention # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] if self.config.is_decoder and encoder_hidden_states is not None: encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) if encoder_attention_mask is None: encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) else: encoder_extended_attention_mask = None # Prepare head mask if needed # 1.0 in head_mask indicate we keep the head # attention_probs has shape bsz x n_heads x N x N # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) encoder_outputs = self.encoder( h_input, attention_mask=extended_attention_mask, head_mask=head_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_extended_attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = encoder_outputs[0] pooled_output = self.pooler(sequence_output) if self.pooler is not None else None if not return_dict: return (sequence_output, pooled_output) + encoder_outputs[1:] return BaseModelOutputWithPoolingAndCrossAttentions( last_hidden_state=sequence_output, pooler_output=pooled_output, past_key_values=encoder_outputs.past_key_values, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, cross_attentions=encoder_outputs.cross_attentions, ) def get_extended_attention_mask( self, attention_mask: Tensor, input_shape: Tuple[int], device: device = None, dtype: torch.float = None ) -> Tensor: """ Makes broadcastable attention and causal masks so that future and masked tokens are ignored. Arguments: attention_mask (`torch.Tensor`): Mask with ones indicating tokens to attend to, zeros for tokens to ignore. input_shape (`Tuple[int]`): The shape of the input to the model. Returns: `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. """ if dtype is None: dtype = torch.float32 if not (attention_mask.dim() == 2 and self.config.is_decoder): # show warning only if it won't be shown in `create_extended_attention_mask_for_decoder` if device is not None: warnings.warn( "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning ) # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] # ourselves in which case we just need to make it broadcastable to all heads. if attention_mask.dim() == 3: extended_attention_mask = attention_mask[:, None, :, :] elif attention_mask.dim() == 2: # Provided a padding mask of dimensions [batch_size, seq_length] # - if the model is a decoder, apply a causal mask in addition to the padding mask # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length] if self.config.is_decoder: extended_attention_mask = ModuleUtilsMixin.create_extended_attention_mask_for_decoder( input_shape, attention_mask, device ) else: extended_attention_mask = attention_mask[:, None, None, :] else: raise ValueError( f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})" ) # Since attention_mask is 1.0 for positions we want to attend and 0.0 for # masked positions, this operation will create a tensor which is 0.0 for # positions we want to attend and the dtype's smallest value for masked positions. # Since we are adding it to the raw scores before the softmax, this is # effectively the same as removing these entirely. extended_attention_mask = extended_attention_mask.to(dtype=dtype) # fp16 compatibility extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min return extended_attention_mask def get_head_mask( self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False ) -> Tensor: """ Prepare the head mask if needed. Args: head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*): The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard). num_hidden_layers (`int`): The number of hidden layers in the model. is_attention_chunked: (`bool`, *optional*, defaults to `False`): Whether or not the attentions scores are computed by chunks or not. Returns: `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with `[None]` for each layer. """ if head_mask is not None: head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) if is_attention_chunked is True: head_mask = head_mask.unsqueeze(-1) else: head_mask = [None] * num_hidden_layers return head_mask def _reset_parameters(self): r"""Initiate parameters in the transformer model.""" for p in self.parameters(): if p.dim() > 1: normal_(p, mean=0.0, std=self.config.initializer_range) def save_weights(self, path): torch.save(self.state_dict(), path) def load_weights(self, path): self.load_state_dict(torch.load(path)) class TAAS(PreTrainedModel): def __init__(self, config, return_last_hidden_state=False): super(TAAS, self).__init__(config) """ :param d_model: d_k = d_v = d_model/nhead = 64, 模型中向量的维度，论文默认值为 512 :param nhead: 多头注意力机制中多头的数量，论文默认为值 8 :param num_encoder_layers: encoder堆叠的数量，也就是论文中的N，论文默认值为6 :param num_decoder_layers: decoder堆叠的数量，也就是论文中的N，论文默认值为6 :param dim_feedforward: 全连接中向量的维度，论文默认值为 2048 :param dropout: 丢弃率，论文中的默认值为 0.1 """ self.config = deepcopy(config) self.return_last_hidden_state = return_last_hidden_state self.dropout = nn.Dropout(self.config.hidden_dropout_prob) # ================ StellarEmbedding ===================== self.embedding = StellarEmbedding(self.config) self.embedding_weights = Parameter(torch.ones(1, 1, self.config.hidden_size)) # ================ StellarModel ===================== self.stellar_config = deepcopy(config) self.stellar_model = StellarModel(self.stellar_config) # ================ TranSAGE ===================== # self.transage_layer = TranSAGE() self.graphormer = Graphormer3D() # ================ 解码部分 ===================== self.encoder_config = deepcopy(config) self.encoder_config.num_hidden_layers = 1 self.encoder = StellarModel(self.encoder_config) self.encoder_out_dim = self.encoder_config.hidden_size # ================ GC任务部分 ===================== self.gc_trans = nn.Linear(self.encoder_out_dim, 16 * 33, bias=True) # ================ MLM任务部分 ===================== self.cls = ErnieForMaskedLM(self.stellar_config).cls # ================ alias任务部分 ===================== self.down_hidden_dim = 512 self.down_kernel_num = 128 self.alias_trans = nn.Linear(self.encoder_out_dim, self.down_hidden_dim, bias=True) self.alias_trans2 = torch.nn.Conv2d(1, self.down_kernel_num, (2, self.down_hidden_dim), stride=1, bias=True) self.alias_layer = nn.Linear(self.down_kernel_num * 5, 2 * 5, bias=True) # ================ AOI任务部分 ===================== self.aoi_trans = nn.Linear(self.encoder_out_dim, self.down_hidden_dim, bias=True) self.aoi_trans2 = torch.nn.Conv2d(1, self.down_kernel_num, (2, self.down_hidden_dim), stride=1, bias=True) self.aoi_layer = nn.Linear(self.down_kernel_num * 5, 2 * 5, bias=True) # ================ HTC任务部分 ===================== self.htc_trans = nn.Linear(self.encoder_out_dim, 5 * 100, bias=True) # ================ NER任务部分 ===================== # self.ner_model = torch.load('ner.pth') self.ner_model = NER_model(vocab_size=11) # self.ner_model.load_state_dict(torch.load('ner.pth')) def forward(self, input_ids, attention_mask, token_type_ids, node_position_ids, spatial_pos, in_degree, out_degree, edge_type_matrix, edge_input, prov_city_mask: Optional[torch.Tensor] = None, sequence_len=6, labels: Optional[torch.Tensor] = None ): """ :param input_ids: [sequence_len * batch_size, src_len] :param attention_mask: [sequence_len * batch_size, src_len] :param token_type_ids: [sequence_len * batch_size, src_len] :param sequence_len: int :param labels: :param is_eval: bool :return: """ batch_size_input = int(input_ids.shape[0] / sequence_len) embedding_output = self.embedding(input_ids=input_ids, token_type_ids=token_type_ids) stellar_predictions = self.stellar_model(embedding_output, input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) last_hidden_state = stellar_predictions[0].contiguous().view(batch_size_input, sequence_len, -1, self.encoder_out_dim) pooler_output = stellar_predictions[1].contiguous().view(batch_size_input, sequence_len, self.encoder_out_dim) h_ = self.graphormer(pooler_output, spatial_pos, in_degree, out_degree, edge_type_matrix, edge_input, node_position_ids) h_ = h_.unsqueeze(2) new_hidden_state = torch.cat((h_, last_hidden_state[:, :, 1:, :]), dim=2) new_hidden_state = new_hidden_state.contiguous().view(batch_size_input * sequence_len, -1, self.encoder_out_dim) encoder_outputs = self.encoder(new_hidden_state, input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) final_hidden_state = encoder_outputs[0] final_pooler_output = encoder_outputs[1].contiguous().view(batch_size_input, sequence_len, self.encoder_out_dim) prediction_scores = self.cls(final_hidden_state) # 用于 MLM 任务 gc_layer_out = self.gc_trans(final_pooler_output) gc_layer_out = gc_layer_out.contiguous().view(-1, 16) htc_layer_out = self.htc_trans(final_pooler_output) htc_layer_out = htc_layer_out.contiguous().view(-1, 100) # MLM loss if labels is not None: # masked_lm_loss = None loss_fct = CrossEntropyLoss() # -100 index = padding token masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) return [gc_layer_out, masked_lm_loss, prediction_scores, htc_layer_out] if self.return_last_hidden_state: return final_pooler_output, pooler_output return gc_layer_out, final_pooler_output, final_hidden_state, prediction_scores, last_hidden_state, htc_layer_out def get_htc_code(self, htc_layer_out): htc_loss_fct = HTCLoss(device=self.device, reduction='mean') htc_pred = htc_loss_fct.get_htc_code(htc_layer_out) return htc_pred def decode_htc_code_2_chn(self, htc_pred): arr = htc_pred with open(remap_code_2_chn_file_path, 'rb') as fr: remap_code_2_chn = pickle.loads(fr.read()) return remap_code_2_chn['{:02d}{:02d}{:02d}{:01d}{:02d}'.format(arr[0], arr[1], arr[2], arr[3], arr[4])] # Address Standarization def addr_standardize(self, address): tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh') encoded_input = tokenizer(address, return_tensors='pt', padding='max_length', truncation=True, # 超过最大长度截断 max_length=60, add_special_tokens=True).to(self.device) word_ids = encoded_input['input_ids'] attention_mask = encoded_input['attention_mask'] length = len(word_ids) node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device) spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device) in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device) edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device) logits = self.ner_model(**encoded_input, node_position_ids = node_position_ids, spatial_pos = spatial_pos, in_degree = in_degree, out_degree = out_degree, edge_type_matrix = edge_type_matrix, edge_input = edge_input,)[0] output = [] ner_labels = torch.argmax(logits, dim=-1) if len(address) == 1: ner_labels = ner_labels.unsqueeze(0) for i in range(len(address)): ner_label = ner_labels[i] word_id = word_ids[i] # cut padding idx = torch.where(attention_mask[i]>0) ner_label = ner_label[idx][1:-1] word_id = word_id[idx][1:-1] # cut other info idx1 = torch.where(ner_label != 0) ner_label = ner_label[idx1].tolist() word_id = word_id[idx1].tolist() # add house info if 8 in ner_label: idx2 = ''.join([str(i) for i in ner_label]).rfind('8') word_id.insert(idx2+1, 2770) ner_label.insert(idx2+1, 8) if 9 in ner_label: idx2 = ''.join([str(i) for i in ner_label]).rfind('9') word_id.insert(idx2+1, 269) word_id.insert(idx2+2, 183) ner_label.insert(idx2+1, 9) ner_label.insert(idx2+2, 9) if 10 in ner_label: idx2 = ''.join([str(i) for i in ner_label]).rfind('10') word_id.insert(idx2+1, 485) ner_label.insert(idx2+1, 10) output.append(tokenizer.decode(word_id).replace(' ', '')) return output # Address Entity Tokenization def addr_entity(self, address): tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh') encoded_input = tokenizer(address, return_tensors='pt', padding='max_length', truncation=True, # 超过最大长度截断 max_length=60, add_special_tokens=True).to(self.device) word_ids = encoded_input['input_ids'] attention_mask = encoded_input['attention_mask'] length = len(word_ids) node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device) spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device) in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device) edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device) logits = self.ner_model(**encoded_input, node_position_ids = node_position_ids, spatial_pos = spatial_pos, in_degree = in_degree, out_degree = out_degree, edge_type_matrix = edge_type_matrix, edge_input = edge_input,)[0] ner_labels = torch.argmax(logits, dim=-1) if len(address) == 1: ner_labels = ner_labels.unsqueeze(0) output = [] tmp = {1:'省', 2:'市', 3:'区', 4:'街道/镇', 5:'道路', 6:'道路号', 7:'poi', 8:'楼栋号', 9:'单元号', 10:'门牌号'} for i in range(len(address)): ner_label = ner_labels[i] word_id = word_ids[i] idx = torch.where(attention_mask[i]>0) ner_label = ner_label[idx][1:-1] word_id = word_id[idx][1:-1] addr_dict = {} addr_dict = dict.fromkeys(tmp.values(),'无') for j in range(1,11): idx = torch.where(ner_label == j) addr_dict[tmp[j]] = ''.join(tokenizer.decode(word_id[idx]).replace(' ','')) output.append(deepcopy(addr_dict)) return output # House Info Extraction def house_info(self, address): tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh') encoded_input = tokenizer(address, return_tensors='pt', padding='max_length', truncation=True, # 超过最大长度截断 max_length=60, add_special_tokens=True).to(self.device) word_ids = encoded_input['input_ids'] attention_mask = encoded_input['attention_mask'] length = len(word_ids) node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device) spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device) in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device) edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device) logits = self.ner_model(**encoded_input, node_position_ids = node_position_ids, spatial_pos = spatial_pos, in_degree = in_degree, out_degree = out_degree, edge_type_matrix = edge_type_matrix, edge_input = edge_input,)[0] ner_labels = torch.argmax(logits, dim=-1) if len(address) == 1: ner_labels = ner_labels.unsqueeze(0) output = [] for i in range(len(address)): ner_label = ner_labels[i] word_id = word_ids[i] idx = torch.where(attention_mask[i]>0) ner_label = ner_label[idx][1:-1] word_id = word_id[idx][1:-1] building = [] unit = [] room = [] for j in range(len(ner_label)): if ner_label[j] == 8: building.append(word_id[j]) elif ner_label[j] == 9: unit.append(word_id[j]) elif ner_label[j] == 10: room.append(word_id[j]) output.append({'楼栋':tokenizer.decode(building).replace(' ',''), '单元':tokenizer.decode(unit).replace(' ',''), '门牌号': tokenizer.decode(room).replace(' ','')}) return output # Address Completion def addr_complet(self, address): tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh') encoded_input = tokenizer(address, return_tensors='pt', padding='max_length', truncation=True, # 超过最大长度截断 max_length=60, add_special_tokens=True).to(self.device) word_ids = encoded_input['input_ids'] attention_mask = encoded_input['attention_mask'] length = len(word_ids) node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device) spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device) in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device) edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device) edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device) logits = self.ner_model(**encoded_input, node_position_ids = node_position_ids, spatial_pos = spatial_pos, in_degree = in_degree, out_degree = out_degree, edge_type_matrix = edge_type_matrix, edge_input = edge_input,)[0] ner_labels = torch.argmax(logits, dim=-1) if len(address) == 1: ner_labels = ner_labels.unsqueeze(0) if isinstance(address, list): address = address[0] # HTC result g2ptl_model = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True) g2ptl_model.eval() g2ptl_output = g2ptl_model(**encoded_input) htc_layer_out = g2ptl_output.htc_layer_out arr = g2ptl_model.get_htc_code(htc_layer_out) htc_pred = '{:02d}{:02d}{:02d}{:01d}{:02d}'.format(arr[0], arr[1], arr[2], arr[3], arr[4]) with open('remap_code_2_chn_with_all_htc.pkl', 'rb') as fr: remap_code_2_chn = pickle.loads(fr.read()) try: htc_list = remap_code_2_chn[htc_pred][-1] except: return address # revise address level of four city if htc_list[0] in ['北京','上海','重庆','天津']: htc_list = htc_list[1:] htc_list.append('') idx = torch.where(attention_mask>0) ner_label = ner_labels[idx][1:-1].cpu().numpy().tolist() word_id = word_ids[idx][1:-1] for i in range(1,5): # judge the lacked address unit if i not in ner_label: if i == 1: address = htc_list[0] + address ner_label = [1] * len(htc_list[0]) + ner_label else : # find the insert position idx = 0 for j in range(len(ner_label)): if ner_label[j] > i: idx = j break address = address[:idx] + htc_list[i-1] + address[idx:] ner_label = ner_label[:idx] + [i] * len(htc_list[i-1]) + ner_label[idx:] return address # Geo-locating from text to geospatial def geolocate(self, address): g2ptl_model = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True) encoded_input = tokenizer(address, return_tensors='pt') g2ptl_model.eval() output = g2ptl_model(**encoded_input) geo_labels = torch.argmax(output.gc_layer_out, dim=-1) output = [s2_label_dict_remap[int(i)] for i in geo_labels] return 's2网格化结果：' + ''.join(output) # Pick-up Estimation Time of Arrival def pickup_ETA(self, address): print('Users can get the address embeddings using model.encode(address) and feed them to your own ETA model.') # Pick-up and Delivery Route Prediction def route_predict(self, route_data): print('Users can get the address embeddings using model.encode(address) and feed them to your own Route Prediction model.') # Address embeddings def encode(self, address): tokenizer = AutoTokenizer.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True) g2ptl_model = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True) encoded_input = tokenizer(address, return_tensors='pt', padding='max_length', truncation=True, # 超过最大长度截断 max_length=60, add_special_tokens=True) g2ptl_model.eval() output = g2ptl_model(**encoded_input) return output.final_hidden_state def _reset_parameters(self): for p in self.parameters(): if p.dim() > 1: xavier_uniform_(p) def generate_square_subsequent_mask(self, sz): mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) return mask # [sz,sz] def save_weights(self, path): torch.save(self.state_dict(), path) def load_weights(self, path): self.load_state_dict(torch.load(path, map_location=torch.device('cpu')), False) def set_pretrained_weights(self, path): pre_train_weights = torch.load(path, map_location=torch.device('cpu')) new_weights = dict() for layer in self.state_dict().keys(): if layer == 'embedding.position_ids': new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_ids'] elif layer == 'embedding.word_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.word_embeddings.weight'] elif layer == 'embedding.position_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_embeddings.weight'] elif layer == 'embedding.token_type_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.token_type_embeddings.weight'] elif layer == 'embedding.task_type_embeddings.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.task_type_embeddings.weight'] elif layer == 'embedding.LayerNorm.weight': new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.weight'] elif layer == 'embedding.LayerNorm.bias': new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.bias'] elif 'stellar_model' in layer: new_weights[layer] = pre_train_weights[layer.replace('stellar_model', 'ernie_model')] elif layer in pre_train_weights.keys(): new_weights[layer] = pre_train_weights[layer] else: new_weights[layer] = self.state_dict()[layer] self.load_state_dict(new_weights)