



'''simple docstring'''





import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
 _lowerCAmelCase										=     argparse.ArgumentParser(
     description=(
         '''Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned'''
         ''' Distillation'''
     )
 )
 parser.add_argument('''--model_type''', default='''bert''', choices=['''bert'''])
 parser.add_argument('''--model_name''', default='''bert-base-uncased''', type=str)
 parser.add_argument('''--dump_checkpoint''', default='''serialization_dir/tf_bert-base-uncased_0247911.pth''', type=str)
 parser.add_argument('''--vocab_transform''', action='''store_true''')
 _lowerCAmelCase										=     parser.parse_args()

 if args.model_type == "bert":
  _lowerCAmelCase										=     BertForMaskedLM.from_pretrained(args.model_name)
  _lowerCAmelCase										=     '''bert'''
 else:
  raise ValueError('''args.model_type should be "bert".''')

 _lowerCAmelCase										=     model.state_dict()
 _lowerCAmelCase										=     {}

 for w in ["word_embeddings", "position_embeddings"]:
  _lowerCAmelCase										=     state_dict[F"""{prefix}.embeddings.{w}.weight"""]
 for w in ["weight", "bias"]:
  _lowerCAmelCase										=     state_dict[F"""{prefix}.embeddings.LayerNorm.{w}"""]

 _lowerCAmelCase										=     0
 for teacher_idx in [0, 2, 4, 7, 9, 11]:
  for w in ["weight", "bias"]:
   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"""
   ]
   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"""
   ]
   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"""
   ]

   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
   ]
   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
   ]

   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"""
   ]
   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"""
   ]
   _lowerCAmelCase										=     state_dict[
       F"""{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"""
   ]
  std_idx += 1

 _lowerCAmelCase										=     state_dict['''cls.predictions.decoder.weight''']
 _lowerCAmelCase										=     state_dict['''cls.predictions.bias''']
 if args.vocab_transform:
  for w in ["weight", "bias"]:
   _lowerCAmelCase										=     state_dict[F"""cls.predictions.transform.dense.{w}"""]
   _lowerCAmelCase										=     state_dict[F"""cls.predictions.transform.LayerNorm.{w}"""]

 print(F"""N layers selected for distillation: {std_idx}""")
 print(F"""Number of params transferred for distillation: {len(compressed_sd.keys())}""")

 print(F"""Save transferred checkpoint to {args.dump_checkpoint}.""")
 torch.save(compressed_sd, args.dump_checkpoint)







'''simple docstring'''





import os
def        _SCREAMING_SNAKE_CASE  (	UpperCamelCase      ):



    """simple docstring"""





    lowerCAmelCase__			:  List[str]											= len(grid[0]      )
    lowerCAmelCase__			:  int											= len(UpperCamelCase      )

    lowerCAmelCase__			:  Optional[int]											= 0
    lowerCAmelCase__			:  Tuple											= 0
    lowerCAmelCase__			:  Optional[Any]											= 0

    # Check vertically, horizontally, diagonally at the same time (only works
    # for nxn grid)
    for i in range(UpperCamelCase      ):
        for j in range(n_rows - 3      ):
            lowerCAmelCase__			:  str											= grid[j][i] * grid[j + 1][i] * grid[j + 2][i] * grid[j + 3][i]
            lowerCAmelCase__			:  Optional[int]											= grid[i][j] * grid[i][j + 1] * grid[i][j + 2] * grid[i][j + 3]

            # Left-to-right diagonal (\) product
            if i < n_columns - 3:
                lowerCAmelCase__			:  Optional[int]											= (
                    grid[i][j]
                    * grid[i + 1][j + 1]
                    * grid[i + 2][j + 2]
                    * grid[i + 3][j + 3]
                )

            # Right-to-left diagonal(/) product
            if i > 2:
                lowerCAmelCase__			:  Tuple											= (
                    grid[i][j]
                    * grid[i - 1][j + 1]
                    * grid[i - 2][j + 2]
                    * grid[i - 3][j + 3]
                )

            lowerCAmelCase__			:  Dict											= max(
                UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase							,   UpperCamelCase      )
            if max_product > largest:
                lowerCAmelCase__			:  Any											= max_product

    return largest
def        _SCREAMING_SNAKE_CASE  (	):



    """simple docstring"""





    lowerCAmelCase__			:  List[str]											= []
    with open(os.path.dirname(UpperCamelCase      ) + """/grid.txt"""      ) as file:
        for line in file:
            grid.append(line.strip("""\n"""      ).split(""" """      )      )

    lowerCAmelCase__			:  Dict											= [[int(UpperCamelCase      ) for i in grid[j]] for j in range(len(UpperCamelCase      )      )]

    return largest_product(UpperCamelCase      )


if __name__ == "__main__":
 print(solution())


