






"""simple docstring"""




import argparse
import json

from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
from transformers.utils import logging


logging.set_verbosity_info()

def     lowerCAmelCase__  (  UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__       ):



  '''simple docstring'''
  # Initialise PyTorch model
  _a						: List[Any]    =			json.load(open(UpperCamelCase__  ,     """r"""       )       )
  _a						: Dict    =			GPTNeoConfig(
      hidden_size=config_json["""n_embd"""]  ,     num_layers=config_json["""n_layer"""]  ,     num_heads=config_json["""n_head"""]  ,     attention_types=config_json["""attention_types"""]  ,     max_position_embeddings=config_json["""n_positions"""]  ,     resid_dropout=config_json["""res_dropout"""]  ,     embed_dropout=config_json["""embed_dropout"""]  ,     attention_dropout=config_json["""attn_dropout"""]  ,     )
  print(F"""Building PyTorch model from configuration: {config}"""       )
  _a						: Optional[Any]    =			GPTNeoForCausalLM(UpperCamelCase__       )

  # Load weights from tf checkpoint
  load_tf_weights_in_gpt_neo(UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__       )

  # Save pytorch-model
  print(F"""Save PyTorch model to {pytorch_dump_path}"""       )
  model.save_pretrained(UpperCamelCase__       )


if __name__ == "__main__":
 _snake_case					   =     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
 )
 parser.add_argument(
     '--config_file',
     default=None,
     type=str,
     required=True,
     help=(
         'The config json file corresponding to the pre-trained mesh-tf model. \n'
         'This specifies the model architecture.'
     ),
 )
 parser.add_argument(
     '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
 )
 _snake_case					   =     parser.parse_args()
 convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)









"""simple docstring"""




# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    '--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000'
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/

