





_A				:							int							= (1, 1_53, 3_70, 3_71, 16_34, 24_67_80_51, 1_15_13_22_19_01_87_63_99_25_65_09_55_97_97_39_71_52_24_01)
_A				:							tuple							= (-1_53, -1, 0, 1.2, 2_00, "A", [], {}, None)


def        __snake_case       (						lowerCAmelCase_  )							->	bool:
    if not isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ) or n < 1:
        return False

    # Initialization of sum and number of digits.
    SCREAMING_SNAKE_CASE__       							= 0
    SCREAMING_SNAKE_CASE__       							= 0
    SCREAMING_SNAKE_CASE__       							= n
    # Calculation of digits of the number
    while temp > 0:
        number_of_digits += 1
        temp //= 1_0
    # Dividing number into separate digits and find Armstrong number
    SCREAMING_SNAKE_CASE__       							= n
    while temp > 0:
        SCREAMING_SNAKE_CASE__       							= temp % 1_0
        total += rem**number_of_digits
        temp //= 1_0
    return n == total


def        __snake_case       (						lowerCAmelCase_  )							->	bool:
    if not isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ) or n < 1:
        return False

    # Init a "histogram" of the digits
    SCREAMING_SNAKE_CASE__       							= [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    SCREAMING_SNAKE_CASE__       							= 0
    SCREAMING_SNAKE_CASE__       							= 0
    SCREAMING_SNAKE_CASE__       							= n
    while temp > 0:
        SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= divmod(lowerCAmelCase_   ,       1_0  )
        digit_histogram[rem] += 1
        digit_total += 1

    for cnt, i in zip(lowerCAmelCase_   ,       range(len(lowerCAmelCase_  )  )  ):
        total += cnt * i**digit_total

    return n == total


def        __snake_case       (						lowerCAmelCase_  )							->	bool:
    if not isinstance(lowerCAmelCase_   ,       lowerCAmelCase_  ) or n < 1:
        return False
    SCREAMING_SNAKE_CASE__       							= len(str(lowerCAmelCase_  )  )  # the power that all digits will be raised to
    # check if sum of each digit multiplied expo times is equal to number
    return n == sum(int(lowerCAmelCase_  ) ** expo for i in str(lowerCAmelCase_  )  )


def        __snake_case       (						)							->	Optional[int]:
    SCREAMING_SNAKE_CASE__       							= int(input('''Enter an integer to see if it is an Armstrong number: '''  ).strip()  )
    print(f'''{num} is {"" if armstrong_number(lowerCAmelCase_  ) else "not "}an Armstrong number.'''  )
    print(f'''{num} is {"" if narcissistic_number(lowerCAmelCase_  ) else "not "}an Armstrong number.'''  )
    print(f'''{num} is {"" if pluperfect_number(lowerCAmelCase_  ) else "not "}an Armstrong number.'''  )


if __name__ == "__main__":
  import doctest

  doctest.testmod()
  main()







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
  _A				:							Any							= Args()
  _A				:							Union[str, Any]							= wandb.init(project="""bigbird-natural-questions""", config=args.__dict__)
  _A				:							Optional[Any]							= dict(logger.config)
  del wandb_args["batch_size"]
  _A				:							Any							= replace(args, **wandb_args)
  _A				:							Dict							= args.base_dir + """-""" + wandb.run.id
  _A				:							Dict							= replace(args, base_dir=base_dir)
  print(args)

  _A				:							Tuple							= load_dataset("""json""", data_files=args.tr_data_path)["""train"""]
  _A				:							Optional[Any]							= load_dataset("""json""", data_files=args.val_data_path)["""train"""]

  # drop extra batch for now
  _A				:							Dict							= range(len(tr_dataset) - len(tr_dataset) % args.batch_size)
  _A				:							int							= tr_dataset.shuffle().select(indices)
  _A				:							Optional[int]							= range(len(val_dataset) - len(val_dataset) % args.batch_size)
  _A				:							Dict							= val_dataset.shuffle().select(indices)

  if os.environ.get("""TRAIN_ON_SMALL""", """false""") == "true":
    _A				:							Optional[Any]							= tr_dataset.shuffle().select(range(8_00_00))
    _A				:							Tuple							= val_dataset.shuffle().select(range(80_00))

  print(tr_dataset)
  print(val_dataset)

  _A				:							int							= FlaxBigBirdForNaturalQuestions.from_pretrained(
      args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
  )
  _A				:							Any							= BigBirdTokenizerFast.from_pretrained(args.model_id)
  _A				:							Optional[Any]							= DataCollator(pad_id=tokenizer.pad_token_id, max_length=40_96)

  _A				:							int							= {
      """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,
  }
  _A,					_A				:							Union[str, Any]							= build_tx(**tx_args)

  _A				:							str							= 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,
  )

  _A				:							Tuple							= None
  _A				:							Dict							= 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`""")
  _A				:							Tuple							= jax_utils.unreplicate(state.params)
  model.save_pretrained(os.path.join(args.base_dir, """final-weights"""), params=params)
