import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#   - FSDP
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_SCREAMING_SNAKE_CASE   		=			1_6
_SCREAMING_SNAKE_CASE   		=			3_2





def  lowercase(  UpperCamelCase_					)		->  Any:







 '''simple docstring'''
 return int(x / 2**20					)
class SCREAMING_SNAKE_CASE_		:







   def __enter__(    self	:	Optional[int]							):





    """simple docstring"""


    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
    UpperCamelCase      =			torch.cuda.memory_allocated()
    UpperCamelCase      =			psutil.Process()

    UpperCamelCase      =			self.cpu_mem_used()
    UpperCamelCase      =			True
    UpperCamelCase      =			threading.Thread(target=self.peak_monitor_func							)
    UpperCamelCase      =			True
    peak_monitor_thread.start()
    return self







   def 	lowerCamelCase_  (    self	:	List[str]							):





    """simple docstring"""


    return self.process.memory_info().rss







   def 	lowerCamelCase_  (    self	:	int							):





    """simple docstring"""


    UpperCamelCase      =			-1

    while True:
     UpperCamelCase      =			max(self.cpu_mem_used()      ,      self.cpu_peak							)

     # can't sleep or will not catch the peak right (this comment is here on purpose)
     # time.sleep(0.001) # 1msec

     if not self.peak_monitoring:
      break







   def __exit__(    self	:	Tuple      ,      *lowerCamelCase_	:	str							):





    """simple docstring"""


    UpperCamelCase      =			False

    gc.collect()
    torch.cuda.empty_cache()
    UpperCamelCase      =			torch.cuda.memory_allocated()
    UpperCamelCase      =			torch.cuda.max_memory_allocated()
    UpperCamelCase      =			bamb(self.end - self.begin							)
    UpperCamelCase      =			bamb(self.peak - self.begin							)

    UpperCamelCase      =			self.cpu_mem_used()
    UpperCamelCase      =			bamb(self.cpu_end - self.cpu_begin							)
    UpperCamelCase      =			bamb(self.cpu_peak - self.cpu_begin							)
    # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   _SCREAMING_SNAKE_CASE   		=			mocked_dataloaders  # noqa: F811





def  lowercase(  UpperCamelCase_       ,				UpperCamelCase_					)		->  Dict:







 '''simple docstring'''


 if os.environ.get("""TESTING_MOCKED_DATALOADERS"""       ,				SCREAMING_SNAKE_CASE__					) == "1":
  UpperCamelCase      =			2

 # New Code #
 # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
 UpperCamelCase      =			FullyShardedDataParallelPlugin(
     state_dict_config=FullStateDictConfig(offload_to_cpu=SCREAMING_SNAKE_CASE__       ,				ranka_only=SCREAMING_SNAKE_CASE__					)       ,				optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=SCREAMING_SNAKE_CASE__       ,				ranka_only=SCREAMING_SNAKE_CASE__					)       ,				)

 # Initialize accelerator
 if args.with_tracking:
  UpperCamelCase      =			Accelerator(
      cpu=args.cpu       ,				mixed_precision=args.mixed_precision       ,				log_with="""wandb"""       ,				project_dir=args.logging_dir       ,				fsdp_plugin=SCREAMING_SNAKE_CASE__       ,				)
 else:
  UpperCamelCase      =			Accelerator(fsdp_plugin=SCREAMING_SNAKE_CASE__					)
 accelerator.print(accelerator.distributed_type					)

 if hasattr(args.checkpointing_steps       ,				"""isdigit"""					):
  if args.checkpointing_steps == "epoch":
   UpperCamelCase      =			args.checkpointing_steps
  elif args.checkpointing_steps.isdigit():
   UpperCamelCase      =			int(args.checkpointing_steps					)
  else:
   raise ValueError(
       f"""Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."""					)
 else:
  UpperCamelCase      =			None
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 UpperCamelCase      =			config['''lr''']
 UpperCamelCase      =			int(config["""num_epochs"""]					)
 UpperCamelCase      =			int(config["""seed"""]					)
 UpperCamelCase      =			int(config["""batch_size"""]					)

 # We need to initialize the trackers we use, and also store our configuration
 if args.with_tracking:
  UpperCamelCase      =			vars(SCREAMING_SNAKE_CASE__					)
  accelerator.init_trackers("""fsdp_glue_no_trainer"""       ,				SCREAMING_SNAKE_CASE__					)

 UpperCamelCase      =			AutoTokenizer.from_pretrained(args.model_name_or_path					)
 UpperCamelCase      =			load_dataset("""glue"""       ,				"""mrpc"""					)
 UpperCamelCase      =			evaluate.load("""glue"""       ,				"""mrpc"""					)

 def tokenize_function(UpperCamelCase_					):
  # max_length=None => use the model max length (it's actually the default)
  UpperCamelCase      =			tokenizer(examples["""sentence1"""]       ,				examples["""sentence2"""]       ,				truncation=SCREAMING_SNAKE_CASE__       ,				max_length=SCREAMING_SNAKE_CASE__					)
  return outputs

 # Apply the method we just defined to all the examples in all the splits of the dataset
 # starting with the main process first:
 with accelerator.main_process_first():
  UpperCamelCase      =			datasets.map(
      SCREAMING_SNAKE_CASE__       ,				batched=SCREAMING_SNAKE_CASE__       ,				remove_columns=["""idx""", """sentence1""", """sentence2"""]       ,				)

 # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
 # transformers library
 UpperCamelCase      =			tokenized_datasets.rename_column("""label"""       ,				"""labels"""					)

 # If the batch size is too big we use gradient accumulation
 UpperCamelCase      =			1
 if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
  UpperCamelCase      =			batch_size // MAX_GPU_BATCH_SIZE
  UpperCamelCase      =			MAX_GPU_BATCH_SIZE

 def collate_fn(UpperCamelCase_					):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  UpperCamelCase      =			128 if accelerator.distributed_type == DistributedType.TPU else None
  # When using mixed precision we want round multiples of 8/16
  if accelerator.mixed_precision == "fp8":
   UpperCamelCase      =			16
  elif accelerator.mixed_precision != "no":
   UpperCamelCase      =			8
  else:
   UpperCamelCase      =			None

  return tokenizer.pad(
      SCREAMING_SNAKE_CASE__       ,				padding="""longest"""       ,				max_length=SCREAMING_SNAKE_CASE__       ,				pad_to_multiple_of=SCREAMING_SNAKE_CASE__       ,				return_tensors="""pt"""       ,				)

 # Instantiate dataloaders.
 UpperCamelCase      =			DataLoader(
     tokenized_datasets["""train"""]       ,				shuffle=SCREAMING_SNAKE_CASE__       ,				collate_fn=SCREAMING_SNAKE_CASE__       ,				batch_size=SCREAMING_SNAKE_CASE__					)
 UpperCamelCase      =			DataLoader(
     tokenized_datasets["""validation"""]       ,				shuffle=SCREAMING_SNAKE_CASE__       ,				collate_fn=SCREAMING_SNAKE_CASE__       ,				batch_size=SCREAMING_SNAKE_CASE__					)

 set_seed(SCREAMING_SNAKE_CASE__					)

 # Instantiate the model (we build the model here so that the seed also control new weights initialization)
 UpperCamelCase      =			AutoModelForSequenceClassification.from_pretrained(
     args.model_name_or_path       ,				return_dict=SCREAMING_SNAKE_CASE__       ,				low_cpu_mem_usage=SCREAMING_SNAKE_CASE__					)

 # New Code #
 # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
 UpperCamelCase      =			accelerator.prepare(SCREAMING_SNAKE_CASE__					)
 accelerator.print(SCREAMING_SNAKE_CASE__					)

 # Instantiate optimizer
 # New Code #
 # For FSDP feature, at present it doesn't support multiple parameter groups,
 # so we need to create a single parameter group for the whole model
 UpperCamelCase      =			torch.optim.AdamW(params=model.parameters()       ,				lr=SCREAMING_SNAKE_CASE__       ,				weight_decay=2E-4					)

 # Instantiate scheduler
 UpperCamelCase      =			get_linear_schedule_with_warmup(
     optimizer=SCREAMING_SNAKE_CASE__       ,				num_warmup_steps=10       ,				num_training_steps=(len(SCREAMING_SNAKE_CASE__					) * num_epochs) // gradient_accumulation_steps       ,				)

 # New Code #
 # For FSDP feature, prepare everything except the model as we have already prepared the model
 # before creating the optimizer
 # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
 # prepare method.
 UpperCamelCase      =			accelerator.prepare(
     SCREAMING_SNAKE_CASE__       ,				SCREAMING_SNAKE_CASE__       ,				SCREAMING_SNAKE_CASE__       ,				SCREAMING_SNAKE_CASE__					)

 UpperCamelCase      =			0

 # Potentially load in the weights and states from a previous save
 if args.resume_from_checkpoint:
  if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
   accelerator.print(f"""Resumed from checkpoint: {args.resume_from_checkpoint}"""					)
   accelerator.load_state(args.resume_from_checkpoint					)
   UpperCamelCase      =			os.path.basename(args.resume_from_checkpoint					)
  else:
   # Get the most recent checkpoint
   UpperCamelCase      =			[f.name for f in os.scandir(os.getcwd()					) if f.is_dir()]
   dirs.sort(key=os.path.getctime					)
   UpperCamelCase      =			dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
  # Extract `epoch_{i}` or `step_{i}`
  UpperCamelCase      =			os.path.splitext(SCREAMING_SNAKE_CASE__					)[0]

  if "epoch" in training_difference:
   num_epochs -= int(training_difference.replace("""epoch_"""       ,				""""""					)					)
   UpperCamelCase      =			None
  else:
   UpperCamelCase      =			int(training_difference.replace("""step_"""       ,				""""""					)					)
   num_epochs -= resume_step // len(SCREAMING_SNAKE_CASE__					)
   # If resuming by step, we also need to know exactly how far into the DataLoader we went
   UpperCamelCase      =			(num_epochs * len(SCREAMING_SNAKE_CASE__					)) - resume_step

    # Now we train the model
 for epoch in range(SCREAMING_SNAKE_CASE__					):
  # New Code #
  # context manager to track the peak memory usage during the training epoch
  with TorchTracemalloc() as tracemalloc:
   model.train()
   if args.with_tracking:
    UpperCamelCase      =			0
   for step, batch in enumerate(SCREAMING_SNAKE_CASE__					):
    # We need to skip steps until we reach the resumed step
    if args.resume_from_checkpoint and epoch == 0:
     if resume_step is not None and step < resume_step:
      pass
                # We could avoid this line since we set the accelerator with `device_placement=True`.
    batch.to(accelerator.device					)
    UpperCamelCase      =			model(**SCREAMING_SNAKE_CASE__					)
    UpperCamelCase      =			outputs.loss
    UpperCamelCase      =			loss / gradient_accumulation_steps
    # We keep track of the loss at each epoch
    if args.with_tracking:
     total_loss += loss.detach().float()
    accelerator.backward(SCREAMING_SNAKE_CASE__					)
    if step % gradient_accumulation_steps == 0:
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()
     # accelerator.print(lr_scheduler.get_lr())

    overall_step += 1

    if isinstance(SCREAMING_SNAKE_CASE__       ,				SCREAMING_SNAKE_CASE__					):
     UpperCamelCase      =			f"""step_{overall_step}"""
     if overall_step % checkpointing_steps == 0:
      if args.output_dir is not None:
       UpperCamelCase      =			os.path.join(args.output_dir       ,				SCREAMING_SNAKE_CASE__					)
      accelerator.save_state(SCREAMING_SNAKE_CASE__					)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
  accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin					)					)					)
  accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used					)					)
  accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked					)					)
  accelerator.print(
      """Total Peak Memory consumed during the train (max): {}""".format(
          tracemalloc.peaked + bamb(tracemalloc.begin					)					)					)
  # Logging the peak memory usage of the GPU to the tracker
  if args.with_tracking:
   accelerator.log(
       {
           """train_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin					),
       }       ,				step=SCREAMING_SNAKE_CASE__       ,				)

  # New Code #
  # context manager to track the peak memory usage during the evaluation
  with TorchTracemalloc() as tracemalloc:
   model.eval()
   for step, batch in enumerate(SCREAMING_SNAKE_CASE__					):
    # We could avoid this line since we set the accelerator with `device_placement=True`.
    batch.to(accelerator.device					)
    with torch.no_grad():
     UpperCamelCase      =			model(**SCREAMING_SNAKE_CASE__					)
    UpperCamelCase      =			outputs.logits.argmax(dim=-1					)
    UpperCamelCase      =			accelerator.gather_for_metrics((predictions, batch["""labels"""])					)
    metric.add_batch(
        predictions=SCREAMING_SNAKE_CASE__       ,				references=SCREAMING_SNAKE_CASE__       ,				)

   UpperCamelCase      =			metric.compute()
   # Use accelerator.print to print only on the main process.
   accelerator.print(f"""epoch {epoch}:"""       ,				SCREAMING_SNAKE_CASE__					)
   if args.with_tracking:
    accelerator.log(
        {
            """accuracy""": eval_metric["""accuracy"""],
            """f1""": eval_metric["""f1"""],
            """train_loss""": total_loss.item() / len(SCREAMING_SNAKE_CASE__					),
        }       ,				step=SCREAMING_SNAKE_CASE__       ,				)

   if checkpointing_steps == "epoch":
    UpperCamelCase      =			f"""epoch_{epoch}"""
    if args.output_dir is not None:
     UpperCamelCase      =			os.path.join(args.output_dir       ,				SCREAMING_SNAKE_CASE__					)
    accelerator.save_state(SCREAMING_SNAKE_CASE__					)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
  accelerator.print("""Memory before entering the eval : {}""".format(bamb(tracemalloc.begin					)					)					)
  accelerator.print("""Memory consumed at the end of the eval (end-begin): {}""".format(tracemalloc.used					)					)
  accelerator.print("""Peak Memory consumed during the eval (max-begin): {}""".format(tracemalloc.peaked					)					)
  accelerator.print(
      """Total Peak Memory consumed during the eval (max): {}""".format(tracemalloc.peaked + bamb(tracemalloc.begin					)					)					)
  # Logging the peak memory usage of the GPU to the tracker
  if args.with_tracking:
   accelerator.log(
       {
           """eval_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin					),
       }       ,				step=SCREAMING_SNAKE_CASE__       ,				)

 if args.with_tracking:
  accelerator.end_training()





def  lowercase(  )		->  List[str]:







 '''simple docstring'''
 UpperCamelCase      =			argparse.ArgumentParser(description="""Simple example of training script."""					)
 parser.add_argument(
     """--mixed_precision"""       ,				type=SCREAMING_SNAKE_CASE__       ,				default=SCREAMING_SNAKE_CASE__       ,				choices=["""no""", """fp16""", """bf16""", """fp8"""]       ,				help="""Whether to use mixed precision. Choose"""
     """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
     """and an Nvidia Ampere GPU."""       ,				)
 parser.add_argument("""--cpu"""       ,				action="""store_true"""       ,				help="""If passed, will train on the CPU."""					)
 parser.add_argument(
     """--checkpointing_steps"""       ,				type=SCREAMING_SNAKE_CASE__       ,				default=SCREAMING_SNAKE_CASE__       ,				help="""Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch."""       ,				)
 parser.add_argument(
     """--resume_from_checkpoint"""       ,				type=SCREAMING_SNAKE_CASE__       ,				default=SCREAMING_SNAKE_CASE__       ,				help="""If the training should continue from a checkpoint folder."""       ,				)
 parser.add_argument(
     """--with_tracking"""       ,				action="""store_true"""       ,				help="""Whether to load in all available experiment trackers from the environment and use them for logging."""       ,				)
 parser.add_argument(
     """--output_dir"""       ,				type=SCREAMING_SNAKE_CASE__       ,				default="""."""       ,				help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""       ,				)
 parser.add_argument(
     """--logging_dir"""       ,				type=SCREAMING_SNAKE_CASE__       ,				default="""logs"""       ,				help="""Location on where to store experiment tracking logs`"""       ,				)
 parser.add_argument(
     """--model_name_or_path"""       ,				type=SCREAMING_SNAKE_CASE__       ,				help="""Path to pretrained model or model identifier from huggingface.co/models."""       ,				required=SCREAMING_SNAKE_CASE__       ,				)
 UpperCamelCase      =			parser.parse_args()
 UpperCamelCase      =			{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
 training_function(SCREAMING_SNAKE_CASE__       ,				SCREAMING_SNAKE_CASE__					)


if __name__ == "__main__":
   main()

import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         CpmAntConfig,
         CpmAntForCausalLM,
         CpmAntModel,
         CpmAntTokenizer,
     )
@require_torch
class SCREAMING_SNAKE_CASE_		:







   def __init__(    self	:	Optional[int]      ,      lowerCamelCase_	:	List[str]      ,      lowerCamelCase_	:	Tuple=2      ,      lowerCamelCase_	:	List[Any]=8      ,      lowerCamelCase_	:	Any=True      ,      lowerCamelCase_	:	int=False      ,      lowerCamelCase_	:	Optional[Any]=False      ,      lowerCamelCase_	:	Optional[int]=False      ,      lowerCamelCase_	:	Any=False      ,      lowerCamelCase_	:	List[Any]=99      ,      lowerCamelCase_	:	Optional[int]=32      ,      lowerCamelCase_	:	List[str]=3      ,      lowerCamelCase_	:	List[str]=4      ,      lowerCamelCase_	:	Tuple=37      ,      lowerCamelCase_	:	str=32      ,      lowerCamelCase_	:	Any=128      ,      lowerCamelCase_	:	List[str]=8      ,      lowerCamelCase_	:	Union[str, Any]=8      ,      lowerCamelCase_	:	Any=8      ,      lowerCamelCase_	:	Optional[int]=1.0      ,      lowerCamelCase_	:	int=True      ,      ):





    """simple docstring"""


    UpperCamelCase      =			parent
    UpperCamelCase      =			batch_size
    UpperCamelCase      =			seq_length
    UpperCamelCase      =			is_training
    UpperCamelCase      =			use_token_type_ids
    UpperCamelCase      =			use_input_mask
    UpperCamelCase      =			use_labels
    UpperCamelCase      =			use_mc_token_ids
    UpperCamelCase      =			vocab_size
    UpperCamelCase      =			hidden_size
    UpperCamelCase      =			num_hidden_layers
    UpperCamelCase      =			num_attention_heads
    UpperCamelCase      =			intermediate_size
    UpperCamelCase      =			num_buckets
    UpperCamelCase      =			max_distance
    UpperCamelCase      =			prompt_length
    UpperCamelCase      =			prompt_types
    UpperCamelCase      =			segment_types
    UpperCamelCase      =			init_std
    UpperCamelCase      =			return_dict







   def 	lowerCamelCase_  (    self	:	List[str]							):





    """simple docstring"""


    UpperCamelCase      =			{}
    UpperCamelCase      =			ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size							).type(torch.intaa							)
    UpperCamelCase      =			False

    UpperCamelCase      =			self.get_config()

    return (config, input_ids)







   def 	lowerCamelCase_  (    self	:	Optional[Any]							):





    """simple docstring"""


    return CpmAntConfig(
        vocab_size=self.vocab_size      ,      hidden_size=self.hidden_size      ,      num_hidden_layers=self.num_hidden_layers      ,      num_attention_heads=self.num_attention_heads      ,      dim_ff=self.intermediate_size      ,      position_bias_num_buckets=self.num_buckets      ,      position_bias_max_distance=self.max_distance      ,      prompt_types=self.prompt_types      ,      prompt_length=self.prompt_length      ,      segment_types=self.segment_types      ,      use_cache=lowerCamelCase_      ,      init_std=self.init_std      ,      return_dict=self.return_dict      ,      )







   def 	lowerCamelCase_  (    self	:	Optional[int]      ,      lowerCamelCase_	:	List[str]      ,      lowerCamelCase_	:	List[str]      ,      *lowerCamelCase_	:	str							):





    """simple docstring"""


    UpperCamelCase      =			CpmAntModel(config=lowerCamelCase_							)
    model.to(lowerCamelCase_							)
    model.eval()

    UpperCamelCase      =			model(**lowerCamelCase_							).last_hidden_state

    self.parent.assertEqual(hidden_states.shape      ,      (self.batch_size, self.seq_length, config.hidden_size)							)







   def 	lowerCamelCase_  (    self	:	Any      ,      lowerCamelCase_	:	Dict      ,      lowerCamelCase_	:	Optional[Any]      ,      *lowerCamelCase_	:	str							):





    """simple docstring"""


    UpperCamelCase      =			CpmAntForCausalLM(lowerCamelCase_							)
    model.to(lowerCamelCase_							)
    UpperCamelCase      =			input_ids["""input_ids"""].to(lowerCamelCase_							)
    model.eval()

    UpperCamelCase      =			model(**lowerCamelCase_							)
    self.parent.assertEqual(
        model_output.logits.shape      ,      (self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)      ,      )







   def 	lowerCamelCase_  (    self	:	str							):





    """simple docstring"""


    UpperCamelCase     ,    UpperCamelCase      =			self.prepare_config_and_inputs()
    return config, inputs_dict





@require_torch
class SCREAMING_SNAKE_CASE_		(			__lowerCAmelCase     ,    __lowerCAmelCase     ,    unittest.TestCase		):
   __lowerCAmelCase											=							(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
   __lowerCAmelCase											=							(
       {"""feature-extraction""": CpmAntModel, """text-generation""": CpmAntForCausalLM} if is_torch_available() else {}
   )

   __lowerCAmelCase											=							False
   __lowerCAmelCase											=							False
   __lowerCAmelCase											=							False
   __lowerCAmelCase											=							False
   __lowerCAmelCase											=							False







   def 	lowerCamelCase_  (    self	:	List[Any]							):





    """simple docstring"""


    UpperCamelCase      =			CpmAntModelTester(self							)
    UpperCamelCase      =			ConfigTester(self      ,      config_class=lowerCamelCase_							)







   def 	lowerCamelCase_  (    self	:	str							):





    """simple docstring"""


    self.config_tester.create_and_test_config_common_properties()
    self.config_tester.create_and_test_config_to_json_string()
    self.config_tester.create_and_test_config_to_json_file()
    self.config_tester.create_and_test_config_from_and_save_pretrained()
    self.config_tester.check_config_can_be_init_without_params()
    self.config_tester.check_config_arguments_init()







   def 	lowerCamelCase_  (    self	:	int							):





    """simple docstring"""


    unittest.skip("""CPMAnt doesn't support input_embeds."""							)(self.test_inputs_embeds							)







   def 	lowerCamelCase_  (    self	:	Tuple							):





    """simple docstring"""


    unittest.skip(
        """CPMAnt doesn't support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model."""							)(self.test_retain_grad_hidden_states_attentions							)







   def 	lowerCamelCase_  (    self	:	Optional[int]							):





    """simple docstring"""


    UpperCamelCase     ,    UpperCamelCase      =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_cpmant_model(lowerCamelCase_      ,      lowerCamelCase_							)







   def 	lowerCamelCase_  (    self	:	str							):





    """simple docstring"""


    UpperCamelCase     ,    UpperCamelCase      =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_lm_head_model(lowerCamelCase_      ,      lowerCamelCase_							)





@require_torch
class SCREAMING_SNAKE_CASE_		(			unittest.TestCase		):







   @tooslow
   def 	lowerCamelCase_  (    self	:	List[str]							):





    """simple docstring"""


    UpperCamelCase      =			"""今天天气真好！"""
    UpperCamelCase      =			"""openbmb/cpm-ant-10b"""
    UpperCamelCase      =			CpmAntModel.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			CpmAntTokenizer.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			tokenizer(lowerCamelCase_      ,      return_tensors="""pt"""							)
    UpperCamelCase      =			model(**lowerCamelCase_							).last_hidden_state

    UpperCamelCase      =			torch.tensor(
        [[[6.1_7_0_8, 5.9_2_4_4, 1.0_8_3_5], [6.5_2_0_7, 6.2_8_9_3, -1_1.3_3_2_4], [-1.0_1_0_7, -0.0_5_7_6, -5.9_5_7_7]]]      ,      )
    self.assertTrue(torch.allclose(hidden_states[:, :3, :3]      ,      lowerCamelCase_      ,      atol=1E-2							)							)





@require_torch
class SCREAMING_SNAKE_CASE_		(			unittest.TestCase		):







   @tooslow
   def 	lowerCamelCase_  (    self	:	str							):





    """simple docstring"""


    UpperCamelCase      =			"""今天天气真好！"""
    UpperCamelCase      =			"""openbmb/cpm-ant-10b"""
    UpperCamelCase      =			CpmAntForCausalLM.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			CpmAntTokenizer.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			tokenizer(lowerCamelCase_      ,      return_tensors="""pt"""							)
    UpperCamelCase      =			model(**lowerCamelCase_							).logits

    UpperCamelCase      =			torch.tensor(
        [[[-6.4_2_6_7, -6.4_0_8_3, -6.3_9_5_8], [-5.8_8_0_2, -5.9_4_4_7, -5.7_8_1_1], [-5.3_8_9_6, -5.4_8_2_0, -5.4_2_9_5]]]      ,      )
    self.assertTrue(torch.allclose(hidden_states[:, :3, :3]      ,      lowerCamelCase_      ,      atol=1E-2							)							)







   @tooslow
   def 	lowerCamelCase_  (    self	:	Dict							):





    """simple docstring"""


    UpperCamelCase      =			"""openbmb/cpm-ant-10b"""
    UpperCamelCase      =			CpmAntForCausalLM.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			CpmAntTokenizer.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			"""今天天气不错，"""
    UpperCamelCase      =			"""今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的"""
    UpperCamelCase      =			tokenizer(lowerCamelCase_      ,      return_tensors="""pt"""							)
    UpperCamelCase      =			model.generate(**lowerCamelCase_							)
    UpperCamelCase      =			tokenizer.batch_decode(lowerCamelCase_							)
    self.assertEqual(lowerCamelCase_      ,      lowerCamelCase_							)







   @tooslow
   def 	lowerCamelCase_  (    self	:	int							):





    """simple docstring"""


    UpperCamelCase      =			"""openbmb/cpm-ant-10b"""
    UpperCamelCase      =			CpmAntForCausalLM.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			CpmAntTokenizer.from_pretrained(lowerCamelCase_							)
    UpperCamelCase      =			["""今天天气不错，""", """新年快乐，万事如意！"""]
    UpperCamelCase      =			[
        """今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的""",
        """新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的""",
    ]
    UpperCamelCase      =			tokenizer(lowerCamelCase_      ,      return_tensors="""pt"""      ,      padding=lowerCamelCase_							)
    UpperCamelCase      =			model.generate(**lowerCamelCase_							)
    UpperCamelCase      =			tokenizer.batch_decode(lowerCamelCase_							)
    self.assertEqual(lowerCamelCase_      ,      lowerCamelCase_							)
