
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
__magic_name__            =  logging.get_logger(__name__)





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    snake_case__						=				YolosConfig()

    # size of the architecture
    if "yolos_ti" in yolos_name:
        snake_case__						=				192
        snake_case__						=				768
        snake_case__						=				12
        snake_case__						=				3
        snake_case__						=				[800, 1_333]
        snake_case__						=				False
    elif yolos_name == "yolos_s_dWr":
        snake_case__						=				330
        snake_case__						=				14
        snake_case__						=				6
        snake_case__						=				1_320
    elif "yolos_s" in yolos_name:
        snake_case__						=				384
        snake_case__						=				1_536
        snake_case__						=				12
        snake_case__						=				6
    elif "yolos_b" in yolos_name:
        snake_case__						=				[800, 1_344]

    snake_case__						=				91
    snake_case__						=				"huggingface/label-files"
    snake_case__						=				"coco-detection-id2label.json"
    snake_case__						=				json.load(open(hf_hub_download(__UpperCamelCase    ,		__UpperCamelCase    ,		repo_type="dataset"							)    ,		"r"							)							)
    snake_case__						=				{int(__UpperCamelCase							): v for k, v in idalabel.items()}
    snake_case__						=				idalabel
    snake_case__						=				{v: k for k, v in idalabel.items()}

    return config





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase = False							):
    for i in range(config.num_hidden_layers							):
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        snake_case__						=				state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""							)
        snake_case__						=				state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""							)
        # next, add query, keys and values (in that order) to the state dict
        snake_case__						=				in_proj_weight[: config.hidden_size, :]
        snake_case__						=				in_proj_bias[: config.hidden_size]
        snake_case__						=				in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        snake_case__						=				in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        snake_case__						=				in_proj_weight[-config.hidden_size :, :]
        snake_case__						=				in_proj_bias[-config.hidden_size :]





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    if "backbone" in name:
        snake_case__						=				name.replace("backbone"    ,		"vit"							)
    if "cls_token" in name:
        snake_case__						=				name.replace("cls_token"    ,		"embeddings.cls_token"							)
    if "det_token" in name:
        snake_case__						=				name.replace("det_token"    ,		"embeddings.detection_tokens"							)
    if "mid_pos_embed" in name:
        snake_case__						=				name.replace("mid_pos_embed"    ,		"encoder.mid_position_embeddings"							)
    if "pos_embed" in name:
        snake_case__						=				name.replace("pos_embed"    ,		"embeddings.position_embeddings"							)
    if "patch_embed.proj" in name:
        snake_case__						=				name.replace("patch_embed.proj"    ,		"embeddings.patch_embeddings.projection"							)
    if "blocks" in name:
        snake_case__						=				name.replace("blocks"    ,		"encoder.layer"							)
    if "attn.proj" in name:
        snake_case__						=				name.replace("attn.proj"    ,		"attention.output.dense"							)
    if "attn" in name:
        snake_case__						=				name.replace("attn"    ,		"attention.self"							)
    if "norm1" in name:
        snake_case__						=				name.replace("norm1"    ,		"layernorm_before"							)
    if "norm2" in name:
        snake_case__						=				name.replace("norm2"    ,		"layernorm_after"							)
    if "mlp.fc1" in name:
        snake_case__						=				name.replace("mlp.fc1"    ,		"intermediate.dense"							)
    if "mlp.fc2" in name:
        snake_case__						=				name.replace("mlp.fc2"    ,		"output.dense"							)
    if "class_embed" in name:
        snake_case__						=				name.replace("class_embed"    ,		"class_labels_classifier"							)
    if "bbox_embed" in name:
        snake_case__						=				name.replace("bbox_embed"    ,		"bbox_predictor"							)
    if "vit.norm" in name:
        snake_case__						=				name.replace("vit.norm"    ,		"vit.layernorm"							)

    return name





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase							):
    for key in orig_state_dict.copy().keys():
        snake_case__						=				orig_state_dict.pop(__UpperCamelCase							)

        if "qkv" in key:
            snake_case__						=				key.split("."							)
            snake_case__						=				int(key_split[2]							)
            snake_case__						=				model.vit.encoder.layer[layer_num].attention.attention.all_head_size
            if "weight" in key:
                snake_case__						=				val[:dim, :]
                snake_case__						=				val[
                    dim : dim * 2, :
                ]
                snake_case__						=				val[-dim:, :]
            else:
                snake_case__						=				val[:dim]
                snake_case__						=				val[dim : dim * 2]
                snake_case__						=				val[-dim:]
        else:
            snake_case__						=				val

    return orig_state_dict





def 	SCREAMING_SNAKE_CASE__    (      ):
    snake_case__						=				"http://images.cocodataset.org/val2017/000000039769.jpg"
    snake_case__						=				Image.open(requests.get(__UpperCamelCase    ,		stream=__UpperCamelCase							).raw							)
    return im





@torch.no_grad()
def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase = False							):
    snake_case__						=				get_yolos_config(__UpperCamelCase							)

    # load original state_dict
    snake_case__						=				torch.load(__UpperCamelCase    ,		map_location="cpu"							)["model"]

    # load 🤗 model
    snake_case__						=				YolosForObjectDetection(__UpperCamelCase							)
    model.eval()
    snake_case__						=				convert_state_dict(__UpperCamelCase    ,		__UpperCamelCase							)
    model.load_state_dict(__UpperCamelCase							)

    # Check outputs on an image, prepared by YolosImageProcessor
    snake_case__						=				800 if yolos_name != "yolos_ti" else 512
    snake_case__						=				YolosImageProcessor(format="coco_detection"    ,		size=__UpperCamelCase							)
    snake_case__						=				image_processor(images=prepare_img()    ,		return_tensors="pt"							)
    snake_case__						=				model(**__UpperCamelCase							)
    snake_case__						,					snake_case__						=				outputs.logits, outputs.pred_boxes

    snake_case__						,					snake_case__						=				None, None
    if yolos_name == "yolos_ti":
        snake_case__						=				torch.tensor(
            [[-39.5_022, -11.9_820, -17.6_888], [-29.9_574, -9.9_769, -17.7_691], [-42.3_281, -20.7_200, -30.6_294]]							)
        snake_case__						=				torch.tensor(
            [[0.4_021, 0.0_836, 0.7_979], [0.0_184, 0.2_609, 0.0_364], [0.1_781, 0.2_004, 0.2_095]]							)
    elif yolos_name == "yolos_s_200_pre":
        snake_case__						=				torch.tensor(
            [[-24.0_248, -10.3_024, -14.8_290], [-42.0_392, -16.8_200, -27.4_334], [-27.2_743, -11.8_154, -18.7_148]]							)
        snake_case__						=				torch.tensor(
            [[0.2_559, 0.5_455, 0.4_706], [0.2_989, 0.7_279, 0.1_875], [0.7_732, 0.4_017, 0.4_462]]							)
    elif yolos_name == "yolos_s_300_pre":
        snake_case__						=				torch.tensor(
            [[-36.2_220, -14.4_385, -23.5_457], [-35.6_970, -14.7_583, -21.3_935], [-31.5_939, -13.6_042, -16.8_049]]							)
        snake_case__						=				torch.tensor(
            [[0.7_614, 0.2_316, 0.4_728], [0.7_168, 0.4_495, 0.3_855], [0.4_996, 0.1_466, 0.9_996]]							)
    elif yolos_name == "yolos_s_dWr":
        snake_case__						=				torch.tensor(
            [[-42.8_668, -24.1_049, -41.1_690], [-34.7_456, -14.1_274, -24.9_194], [-33.7_898, -12.1_946, -25.6_495]]							)
        snake_case__						=				torch.tensor(
            [[0.5_587, 0.2_773, 0.0_605], [0.5_004, 0.3_014, 0.9_994], [0.4_999, 0.1_548, 0.9_994]]							)
    elif yolos_name == "yolos_base":
        snake_case__						=				torch.tensor(
            [[-40.6_064, -24.3_084, -32.6_447], [-55.1_990, -30.7_719, -35.5_877], [-51.4_311, -33.3_507, -35.6_462]]							)
        snake_case__						=				torch.tensor(
            [[0.5_555, 0.2_794, 0.0_655], [0.9_049, 0.2_664, 0.1_894], [0.9_183, 0.1_984, 0.1_635]]							)
    else:
        raise ValueError(F"""Unknown yolos_name: {yolos_name}"""							)

    assert torch.allclose(logits[0, :3, :3]    ,		__UpperCamelCase    ,		atol=1E-4							)
    assert torch.allclose(pred_boxes[0, :3, :3]    ,		__UpperCamelCase    ,		atol=1E-4							)

    Path(__UpperCamelCase							).mkdir(exist_ok=__UpperCamelCase							)
    print(F"""Saving model {yolos_name} to {pytorch_dump_folder_path}"""							)
    model.save_pretrained(__UpperCamelCase							)
    print(F"""Saving image processor to {pytorch_dump_folder_path}"""							)
    image_processor.save_pretrained(__UpperCamelCase							)

    if push_to_hub:
        snake_case__						=				{
            "yolos_ti": "yolos-tiny",
            "yolos_s_200_pre": "yolos-small",
            "yolos_s_300_pre": "yolos-small-300",
            "yolos_s_dWr": "yolos-small-dwr",
            "yolos_base": "yolos-base",
        }

        print("Pushing to the hub..."							)
        snake_case__						=				model_mapping[yolos_name]
        image_processor.push_to_hub(__UpperCamelCase    ,		organization="hustvl"							)
        model.push_to_hub(__UpperCamelCase    ,		organization="hustvl"							)


if __name__ == "__main__":
     __magic_name__            =  argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--yolos_name''',
         default='''yolos_s_200_pre''',
         type=str,
         help=(
             '''Name of the YOLOS model you\'d like to convert. Should be one of \'yolos_ti\', \'yolos_s_200_pre\','''
             ''' \'yolos_s_300_pre\', \'yolos_s_dWr\', \'yolos_base\'.'''
         ),
     )
     parser.add_argument(
         '''--checkpoint_path''', default=None, type=str, help='''Path to the original state dict (.pth file).'''
     )
     parser.add_argument(
         '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
     )
     parser.add_argument(
         '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
     )

     __magic_name__            =  parser.parse_args()
     convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)


import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    snake_case__						=				torch.exp(__lowerCAmelCase							)
    snake_case__						=				torch.sum(__lowerCAmelCase    ,		dim=1							)  # sum of exp(x_i)
    snake_case__						=				torch.sum(x * exp_x    ,		dim=1							)  # sum of x_i * exp(x_i)
    return torch.log(__lowerCAmelCase							) - B / A




class 					_SCREAMING_SNAKE_CASE  (      nn.Module ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__()
          snake_case__						=				config.output_attentions
          snake_case__						=				config.output_hidden_states
          snake_case__						=				nn.ModuleList([BertLayer(lowerCamelCase ) for _ in range(config.num_hidden_layers )] )
          snake_case__						=				nn.ModuleList([BertHighway(lowerCamelCase ) for _ in range(config.num_hidden_layers )] )

          snake_case__						=				[-1 for _ in range(config.num_hidden_layers )]
      def A_			(			self		,	lowerCamelCase ):
          if (type(lowerCamelCase ) is float) or (type(lowerCamelCase ) is int):
              for i in range(len(self.early_exit_entropy ) ):
                  snake_case__						=				x
          else:
              snake_case__						=				x
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				pooler.state_dict()
          for highway in self.highway:
              for name, param in highway.pooler.state_dict().items():
                  param.copy_(loaded_model[name] )



      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	):
          snake_case__						=				()
          snake_case__						=				()
          snake_case__						=				()
          for i, layer_module in enumerate(self.layer ):
              if self.output_hidden_states:
                  snake_case__						=				all_hidden_states + (hidden_states,)

              snake_case__						=				layer_module(
                  lowerCamelCase		,	lowerCamelCase		,	head_mask[i]		,	lowerCamelCase		,	lowerCamelCase )
              snake_case__						=				layer_outputs[0]

              if self.output_attentions:
                  snake_case__						=				all_attentions + (layer_outputs[1],)

              snake_case__						=				(hidden_states,)
              if self.output_hidden_states:
                  snake_case__						=				current_outputs + (all_hidden_states,)
              if self.output_attentions:
                  snake_case__						=				current_outputs + (all_attentions,)

              snake_case__						=				self.highway[i](lowerCamelCase )
              # logits, pooled_output

              if not self.training:
                  snake_case__						=				highway_exit[0]
                  snake_case__						=				entropy(lowerCamelCase )
                  snake_case__						=				highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                  snake_case__						=				all_highway_exits + (highway_exit,)

                  if highway_entropy < self.early_exit_entropy[i]:
                      snake_case__						=				(highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                      raise HighwayException(lowerCamelCase		,	i + 1 )
              else:
                  snake_case__						=				all_highway_exits + (highway_exit,)

        # Add last layer
          if self.output_hidden_states:
              snake_case__						=				all_hidden_states + (hidden_states,)

          snake_case__						=				(hidden_states,)
          if self.output_hidden_states:
              snake_case__						=				outputs + (all_hidden_states,)
          if self.output_attentions:
              snake_case__						=				outputs + (all_attentions,)

          snake_case__						=				outputs + (all_highway_exits,)
          return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits





@add_start_docstrings(
    'The Bert Model transformer with early exiting (DeeBERT). '       ,					__UpperCamelCase       ,					)
class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__(lowerCamelCase )
          snake_case__						=				config

          snake_case__						=				BertEmbeddings(lowerCamelCase )
          snake_case__						=				DeeBertEncoder(lowerCamelCase )
          snake_case__						=				BertPooler(lowerCamelCase )

          self.init_weights()
      def A_			(			self ):
          self.encoder.init_highway_pooler(self.pooler )
      def A_			(			self ):
          return self.embeddings.word_embeddings
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				value
      def A_			(			self		,	lowerCamelCase ):
          for layer, heads in heads_to_prune.items():
              self.encoder.layer[layer].attention.prune_heads(lowerCamelCase )



      @add_start_docstrings_to_model_forward(lowerCamelCase )
      def A_			(			self		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	):
          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:
              snake_case__						=				input_ids.size()
          elif inputs_embeds is not None:
              snake_case__						=				inputs_embeds.size()[:-1]
          else:
              raise ValueError("You have to specify either input_ids or inputs_embeds" )

          snake_case__						=				input_ids.device if input_ids is not None else inputs_embeds.device

          if attention_mask is None:
              snake_case__						=				torch.ones(lowerCamelCase		,	device=lowerCamelCase )
          if encoder_attention_mask is None:
              snake_case__						=				torch.ones(lowerCamelCase		,	device=lowerCamelCase )
          if token_type_ids is None:
              snake_case__						=				torch.zeros(lowerCamelCase		,	dtype=torch.long		,	device=lowerCamelCase )

          # 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.
          snake_case__						=				self.get_extended_attention_mask(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase )

          # If a 2D ou 3D attention mask is provided for the cross-attention
          # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if encoder_attention_mask.dim() == 3:
              snake_case__						=				encoder_attention_mask[:, None, :, :]
          if encoder_attention_mask.dim() == 2:
              snake_case__						=				encoder_attention_mask[:, None, None, :]

          snake_case__						=				encoder_extended_attention_mask.to(
              dtype=next(self.parameters() ).dtype )  # fp16 compatibility
          snake_case__						=				(1.0 - encoder_extended_attention_mask) * -1_0_0_0_0.0

          # 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]
          snake_case__						=				self.get_head_mask(lowerCamelCase		,	self.config.num_hidden_layers )

          snake_case__						=				self.embeddings(
              input_ids=lowerCamelCase		,	position_ids=lowerCamelCase		,	token_type_ids=lowerCamelCase		,	inputs_embeds=lowerCamelCase )
          snake_case__						=				self.encoder(
              lowerCamelCase		,	attention_mask=lowerCamelCase		,	head_mask=lowerCamelCase		,	encoder_hidden_states=lowerCamelCase		,	encoder_attention_mask=lowerCamelCase		,	)
          snake_case__						=				encoder_outputs[0]
          snake_case__						=				self.pooler(lowerCamelCase )

          snake_case__						=				(
              sequence_output,
              pooled_output,
          ) + encoder_outputs[
              1:
          ]  # add hidden_states and attentions if they are here
          return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits





class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase ):
          snake_case__						=				message
          snake_case__						=				exit_layer  # start from 1!





class 					_SCREAMING_SNAKE_CASE  (      nn.Module ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__()
          snake_case__						=				BertPooler(lowerCamelCase )
          snake_case__						=				nn.Dropout(config.hidden_dropout_prob )
          snake_case__						=				nn.Linear(config.hidden_size		,	config.num_labels )



      def A_			(			self		,	lowerCamelCase ):
          # Pooler
          snake_case__						=				encoder_outputs[0]
          snake_case__						=				self.pooler(lowerCamelCase )
          # "return" pooler_output

          # BertModel
          snake_case__						=				(pooler_input, pooler_output) + encoder_outputs[1:]
          # "return" bmodel_output

          # Dropout and classification
          snake_case__						=				bmodel_output[1]

          snake_case__						=				self.dropout(lowerCamelCase )
          snake_case__						=				self.classifier(lowerCamelCase )

          return logits, pooled_output





@add_start_docstrings(
    'Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. '       ,					__UpperCamelCase       ,					)
class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__(lowerCamelCase )
          snake_case__						=				config.num_labels
          snake_case__						=				config.num_hidden_layers

          snake_case__						=				DeeBertModel(lowerCamelCase )
          snake_case__						=				nn.Dropout(config.hidden_dropout_prob )
          snake_case__						=				nn.Linear(config.hidden_size		,	self.config.num_labels )

          self.init_weights()



      @add_start_docstrings_to_model_forward(lowerCamelCase )
      def A_			(			self		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=-1		,	lowerCamelCase=False		,	):
          snake_case__						=				self.num_layers
          try:
              snake_case__						=				self.bert(
                  lowerCamelCase		,	attention_mask=lowerCamelCase		,	token_type_ids=lowerCamelCase		,	position_ids=lowerCamelCase		,	head_mask=lowerCamelCase		,	inputs_embeds=lowerCamelCase		,	)
              # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

              snake_case__						=				outputs[1]

              snake_case__						=				self.dropout(lowerCamelCase )
              snake_case__						=				self.classifier(lowerCamelCase )
              snake_case__						=				(logits,) + outputs[2:]  # add hidden states and attention if they are here
          except HighwayException as e:
              snake_case__						=				e.message
              snake_case__						=				e.exit_layer
              snake_case__						=				outputs[0]

          if not self.training:
              snake_case__						=				entropy(lowerCamelCase )
              snake_case__						=				[]
              snake_case__						=				[]
          if labels is not None:
              if self.num_labels == 1:
                  #  We are doing regression
                  snake_case__						=				MSELoss()
                  snake_case__						=				loss_fct(logits.view(-1 )		,	labels.view(-1 ) )
              else:
                  snake_case__						=				CrossEntropyLoss()
                  snake_case__						=				loss_fct(logits.view(-1		,	self.num_labels )		,	labels.view(-1 ) )

              # work with highway exits
              snake_case__						=				[]
              for highway_exit in outputs[-1]:
                  snake_case__						=				highway_exit[0]
                  if not self.training:
                      highway_logits_all.append(lowerCamelCase )
                      highway_entropy.append(highway_exit[2] )
                  if self.num_labels == 1:
                      #  We are doing regression
                      snake_case__						=				MSELoss()
                      snake_case__						=				loss_fct(highway_logits.view(-1 )		,	labels.view(-1 ) )
                  else:
                      snake_case__						=				CrossEntropyLoss()
                      snake_case__						=				loss_fct(highway_logits.view(-1		,	self.num_labels )		,	labels.view(-1 ) )
                  highway_losses.append(lowerCamelCase )

              if train_highway:
                  snake_case__						=				(sum(highway_losses[:-1] ),) + outputs
                  # exclude the final highway, of course
              else:
                  snake_case__						=				(loss,) + outputs
          if not self.training:
              snake_case__						=				outputs + ((original_entropy, highway_entropy), exit_layer)
              if output_layer >= 0:
                  snake_case__						=				(
                      (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                  )  # use the highway of the last layer

          return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)



