import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


UpperCAmelCase_    		=				'''
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
'''





class   lowerCAmelCase    (       unittest.TestCase      ,		_a      ):



     def       a__	(       self							):
            _A=						load_tool('summarization'							)
            self.tool.setup()
            _A=						load_tool('summarization'					,   remote=lowerCAmelCase__							)



     def       a__	(       self							):
            _A=						self.tool(lowerCAmelCase__							)
            self.assertEqual(
                lowerCAmelCase__					,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'					,   )



     def       a__	(       self							):
            _A=						self.remote_tool(lowerCAmelCase__							)
            self.assertEqual(
                lowerCAmelCase__					,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'					,   )



     def       a__	(       self							):
            _A=						self.tool(text=lowerCAmelCase__							)
            self.assertEqual(
                lowerCAmelCase__					,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'					,   )





     def       a__	(       self							):
            _A=						self.remote_tool(text=lowerCAmelCase__							)
            self.assertEqual(
                lowerCAmelCase__					,   'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'					,   )
import os
from dataclasses import replace

import jax
import wandb
from bigbird_flax import Args, DataCollator, FlaxBigBirdForNaturalQuestions, Trainer, build_tx, train_step, val_step
from datasets import load_dataset
from flax import jax_utils

from transformers import BigBirdTokenizerFast


if __name__ == "__main__":
      print('''#################### AVAILABLE DEVICES ####################''')
      print(jax.devices())
      print('''###########################################################''')

      # setup for wandb sweep
      UpperCAmelCase_    		=				Args()
      UpperCAmelCase_    		=				wandb.init(project='''bigbird-natural-questions''', config=args.__dict__)
      UpperCAmelCase_    		=				dict(logger.config)
      del wandb_args["batch_size"]
      UpperCAmelCase_    		=				replace(args, **wandb_args)
      UpperCAmelCase_    		=				args.base_dir + '''-''' + wandb.run.id
      UpperCAmelCase_    		=				replace(args, base_dir=base_dir)
      print(args)

      UpperCAmelCase_    		=				load_dataset('''json''', data_files=args.tr_data_path)['''train''']
      UpperCAmelCase_    		=				load_dataset('''json''', data_files=args.val_data_path)['''train''']

      # drop extra batch for now
      UpperCAmelCase_    		=				range(len(tr_dataset) - len(tr_dataset) % args.batch_size)
      UpperCAmelCase_    		=				tr_dataset.shuffle().select(indices)
      UpperCAmelCase_    		=				range(len(val_dataset) - len(val_dataset) % args.batch_size)
      UpperCAmelCase_    		=				val_dataset.shuffle().select(indices)

      if os.environ.get('''TRAIN_ON_SMALL''', '''false''') == "true":
            UpperCAmelCase_    		=				tr_dataset.shuffle().select(range(80_000))
            UpperCAmelCase_    		=				val_dataset.shuffle().select(range(8_000))

      print(tr_dataset)
      print(val_dataset)

      UpperCAmelCase_    		=				FlaxBigBirdForNaturalQuestions.from_pretrained(
          args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
      )
      UpperCAmelCase_    		=				BigBirdTokenizerFast.from_pretrained(args.model_id)
      UpperCAmelCase_    		=				DataCollator(pad_id=tokenizer.pad_token_id, max_length=4_096)

      UpperCAmelCase_    		=				{
          '''lr''': args.lr,
          '''init_lr''': args.init_lr,
          '''warmup_steps''': args.warmup_steps,
          '''num_train_steps''': args.max_epochs * (len(tr_dataset) // args.batch_size),
          '''weight_decay''': args.weight_decay,
      }
      UpperCAmelCase_		,     UpperCAmelCase_    		=				build_tx(**tx_args)

      UpperCAmelCase_    		=				Trainer(
          args=args,
          data_collator=data_collator,
          model_save_fn=model.save_pretrained,
          train_step_fn=train_step,
          val_step_fn=val_step,
          logger=logger,
          scheduler_fn=lr,
      )

      UpperCAmelCase_    		=				None
      UpperCAmelCase_    		=				trainer.create_state(model, tx, num_train_steps=tx_args['''num_train_steps'''], ckpt_dir=ckpt_dir)
      try:
            trainer.train(state, tr_dataset, val_dataset)
      except KeyboardInterrupt:
            print('''Oooops; TRAINING STOPPED UNFORTUNATELY''')

      print('''SAVING WEIGHTS IN `final-weights`''')
      UpperCAmelCase_    		=				jax_utils.unreplicate(state.params)
      model.save_pretrained(os.path.join(args.base_dir, '''final-weights'''), params=params)