




import logging
import os
import time
from argparse import Namespace
from pathlib import Path

import datasets
import torch
from accelerate import Accelerator, DistributedType
from accelerate.utils import ProjectConfiguration
from arguments import TrainingArguments
from datasets import load_dataset
from huggingface_hub import Repository
from torch.optim import AdamW
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, get_scheduler, set_seed
class      a__       (   snake_case__ ):






						def __init__(  self						,     _A						,     _A						,     _A=False						,     _A=1_0_2_4						,     _A=1_0_2_4						,     _A=3.6						,     _A=False						,     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			tokenizer
											__lowerCAmelCase							  =			tokenizer.bos_token_id
											__lowerCAmelCase							  =			dataset
											__lowerCAmelCase							  =			seq_length
											__lowerCAmelCase							  =			0
											__lowerCAmelCase							  =			infinite
											__lowerCAmelCase							  =			0
											__lowerCAmelCase							  =			tokenized

											if self.tokenized:
																__lowerCAmelCase							  =			seq_length * num_of_sequences
																__lowerCAmelCase							  =			'''input_ids'''
											else:
																__lowerCAmelCase							  =			seq_length * chars_per_token * num_of_sequences
																__lowerCAmelCase							  =			'''content'''






						def __iter__(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			iter(self.dataset     )
											__lowerCAmelCase							  =			True
											while more_examples:
																__lowerCAmelCase							  =			[], 0
																while True:
																					if buffer_len >= self.max_buffer_size:
																										break
																					try:
																										buffer.append(next(__snake_case     )[self.content_field]     )
																										buffer_len += len(buffer[-1]     )
																					except StopIteration:
																										if self.infinite:
																															__lowerCAmelCase							  =			iter(self.dataset     )
																															self.epoch += 1
																															logger.info(f"""Dataset epoch: {self.epoch}"""     )
																										else:
																															__lowerCAmelCase							  =			False
																															break
																if self.tokenized:
																					__lowerCAmelCase							  =			buffer
																else:
																					__lowerCAmelCase							  =			self.tokenizer(__snake_case						,     truncation=__snake_case     )['''input_ids''']
																__lowerCAmelCase							  =			[]
																for tokenized_input in tokenized_inputs:
																					all_token_ids.extend(tokenized_input + [self.concat_token_id]     )
																for i in range(0						,     len(__snake_case     )						,     self.seq_length     ):
																					__lowerCAmelCase							  =			all_token_ids[i : i + self.seq_length]
																					if len(__snake_case     ) == self.seq_length:
																										self.current_size += 1
																										yield torch.tensor(__snake_case     )






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A=1_0_0_0     ):




											"""simple docstring"""



											return ShufflerIterDataPipe(self						,     buffer_size=__snake_case     )







def       _a (		SCREAMING_SNAKE_CASE_				:							Optional[int]	):
					__lowerCAmelCase							  =			args.model_ckpt.split("/"	)[-1]
					__lowerCAmelCase							  =			logging.getLogger(__name__	)
					__lowerCAmelCase							  =			Path(args.save_dir	) / '''log/'''
					log_dir.mkdir(exist_ok=__lowercase	)
					__lowerCAmelCase							  =			F"""debug_{accelerator.process_index}.log"""
					logging.basicConfig(
					    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s" , datefmt="%m/%d/%Y %H:%M:%S" , level=logging.INFO , handlers=[logging.FileHandler(log_dir / filename	), logging.StreamHandler()] , )
					if accelerator.is_main_process:  # we only want to setup logging once
										accelerator.init_trackers(__lowercase , vars(__lowercase	)	)
										__lowerCAmelCase							  =			accelerator.trackers[0].run.name
										logger.setLevel(logging.INFO	)
										datasets.utils.logging.set_verbosity_info()
										transformers.utils.logging.set_verbosity_info()
					else:
										__lowerCAmelCase							  =			''''''
										logger.setLevel(logging.ERROR	)
										datasets.utils.logging.set_verbosity_error()
										transformers.utils.logging.set_verbosity_error()
					return logger, run_name







def       _a (		SCREAMING_SNAKE_CASE_				:							Optional[int]	):
					__lowerCAmelCase							  =			{'''streaming''': True}
					__lowerCAmelCase							  =			load_dataset(args.dataset_name_train , split="train" , **__lowercase	)
					__lowerCAmelCase							  =			train_data.shuffle(buffer_size=args.shuffle_buffer , seed=args.seed	)
					__lowerCAmelCase							  =			load_dataset(args.dataset_name_valid , split="train" , **__lowercase	)
					__lowerCAmelCase							  =			ConstantLengthDataset(
					    __lowercase , __lowercase , infinite=__lowercase , seq_length=args.seq_length , tokenized=args.tokenized	)
					__lowerCAmelCase							  =			ConstantLengthDataset(
					    __lowercase , __lowercase , infinite=__lowercase , seq_length=args.seq_length , tokenized=args.tokenized	)
					__lowerCAmelCase							  =			train_dataset.shuffle(buffer_size=args.shuffle_buffer	)
					__lowerCAmelCase							  =			DataLoader(__lowercase , batch_size=args.train_batch_size , shuffle=__lowercase	)
					__lowerCAmelCase							  =			DataLoader(__lowercase , batch_size=args.valid_batch_size	)
					return train_dataloader, eval_dataloader







def       _a (		SCREAMING_SNAKE_CASE_				:							List[Any] , SCREAMING_SNAKE_CASE_				:							Optional[Any] , SCREAMING_SNAKE_CASE_				:							List[str]=["bias", "ln_1.weight", "ln_2.weight", "ln_f.weight"]	):
					__lowerCAmelCase							  =			[], []
					for n, p in model.named_parameters():
										if any(nd in n for nd in no_decay	):
															params_without_wd.append(__lowercase	)
										else:
															params_with_wd.append(__lowercase	)
					return [
					    {"params": params_with_wd, "weight_decay": args.weight_decay},
					    {"params": params_without_wd, "weight_decay": 0.0},
					]







def       _a (		SCREAMING_SNAKE_CASE_				:							List[Any] , SCREAMING_SNAKE_CASE_				:							int	):
					logger.info(F"""Step {step}: {metrics}"""	)
					if accelerator.is_main_process:
										accelerator.log(__lowercase , __lowercase	)







def       _a (		SCREAMING_SNAKE_CASE_				:							Optional[int] , SCREAMING_SNAKE_CASE_				:							int , SCREAMING_SNAKE_CASE_				:							Optional[int]	):
					# TFLOPs formula (from Equation 3 in Section 5.1 of https://arxiv.org/pdf/2104.04473.pdf).
					__lowerCAmelCase							  =			accelerator.unwrap_model(__lowercase	).config
					__lowerCAmelCase							  =			4 if args.gradient_checkpointing else 3
					__lowerCAmelCase							  =			args.train_batch_size * accelerator.state.num_processes * args.gradient_accumulation_steps
					__lowerCAmelCase							  =			24 * checkpoint_factor * batch_size * args.seq_length * config_model.n_layer * (config_model.n_embd**2)
					__lowerCAmelCase							  =			factor * (
					    1.0
					    + (args.seq_length / (6.0 * config_model.n_embd))
					    + (tokenizer.vocab_size / (16.0 * config_model.n_layer * config_model.n_embd))
					)
					__lowerCAmelCase							  =			flops_per_iteration / (elapsed_time * accelerator.state.num_processes * (10**12))
					return tflops







def       _a (		SCREAMING_SNAKE_CASE_				:							List[Any]	):
					model.eval()
					__lowerCAmelCase							  =			[]
					for step, batch in enumerate(__lowercase	):
										with torch.no_grad():
															__lowerCAmelCase							  =			model(__lowercase , labels=__lowercase	)
										__lowerCAmelCase							  =			outputs.loss.repeat(args.valid_batch_size	)
										losses.append(accelerator.gather(__lowercase	)	)
										if args.max_eval_steps > 0 and step >= args.max_eval_steps:
															break
					__lowerCAmelCase							  =			torch.cat(__lowercase	)
					__lowerCAmelCase							  =			losses[: eval_dataloader.dataset.current_size].mean()
					try:
										__lowerCAmelCase							  =			torch.exp(__lowercase	)
					except OverflowError:
										__lowerCAmelCase							  =			float("inf"	)
					return loss.item(), perplexity.item()


# Settings
UpperCamelCase__            =     HfArgumentParser(TrainingArguments)
UpperCamelCase__            =     parser.parse_args()

# Accelerator
UpperCamelCase__            =     ProjectConfiguration(project_dir=args.save_dir, logging_dir="""log""")
UpperCamelCase__            =     Accelerator(log_with=["""wandb""", """tensorboard"""], project_config=config)
UpperCamelCase__            =     {str(k): str(v) for k, v in accelerator.state.__dict__.items()}

UpperCamelCase__            =     Namespace(**vars(args), **acc_state)
UpperCamelCase__            =     accelerator.state.num_processes * args.train_batch_size
set_seed(args.seed)

# Clone model repository
if accelerator.is_main_process:
					UpperCamelCase__            =     Repository(args.save_dir, clone_from=args.model_ckpt)

# Logging
UpperCamelCase__            =     setup_logging(args)
logger.info(accelerator.state)

# Checkout new branch on repo
if accelerator.is_main_process:
					hf_repo.git_checkout(run_name, create_branch_ok=True)

# Load model and tokenizer
UpperCamelCase__            =     AutoModelForCausalLM.from_pretrained(args.save_dir)
if args.gradient_checkpointing:
					model.gradient_checkpointing_enable()
UpperCamelCase__            =     AutoTokenizer.from_pretrained(args.save_dir)

# Load dataset and dataloader
UpperCamelCase__            =     create_dataloaders(args)

# Prepare the optimizer and learning rate scheduler
UpperCamelCase__            =     AdamW(get_grouped_params(model, args), lr=args.learning_rate)
UpperCamelCase__            =     get_scheduler(
    name=args.lr_scheduler_type,
    optimizer=optimizer,
    num_warmup_steps=args.num_warmup_steps,
    num_training_steps=args.max_train_steps,
)
accelerator.register_for_checkpointing(lr_scheduler)







def       _a (		):
					return optimizer.param_groups[0]["lr"]


# Prepare everything with our `accelerator`.
UpperCamelCase__            =     accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

# 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(args.save_dir) if f.is_dir() and 'step' in str(f)]
										dirs.sort(key=os.path.getctime)
										UpperCamelCase__            =     dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
					# Extract the step of the checkpoint to continue from there
					UpperCamelCase__            =     os.path.splitext(path)[0]
					UpperCamelCase__            =     int(training_difference.replace("""step_""", """"""))

# Train model
model.train()
UpperCamelCase__            =     0
UpperCamelCase__            =     time.time()
UpperCamelCase__            =     0
for step, batch in enumerate(train_dataloader, start=1):
					if args.resume_from_checkpoint and step < resume_step:
										continue  # we need to skip steps until we reach the resumed step
					UpperCamelCase__            =     model(batch, labels=batch, use_cache=False).loss
					UpperCamelCase__            =     accelerator.gather(loss.repeat(args.train_batch_size)).mean()
					loss_tracking += avg_loss.item() / args.gradient_accumulation_steps
					log_metrics(step, {"""samples""": step * samples_per_step, """loss_per_step/train""": loss.item()})
					UpperCamelCase__            =     loss / args.gradient_accumulation_steps
					if step % args.gradient_accumulation_steps != 0:
										# Prevent backward from doing gradient all_reduce in every step
										if accelerator.distributed_type == DistributedType.MULTI_GPU:
															with model.no_sync():
																				accelerator.backward(loss)
										else:
															accelerator.backward(loss)
					else:
										UpperCamelCase__            =     get_lr()
										accelerator.backward(loss)
										accelerator.clip_grad_norm_(model.parameters(), 1.0)
										optimizer.step()
										lr_scheduler.step()
										optimizer.zero_grad()
										UpperCamelCase__            =     time.time() - t_start
										UpperCamelCase__            =     compute_tflops(elapsed_time, accelerator, args)
										log_metrics(
										    step,
										    {
										        """steps""": completed_steps,
										        """loss/train""": loss_tracking,
										        """lr""": lr,
										        """tflops""": tflops,
										        """time_per_iteration""": elapsed_time,
										    },
										)
										UpperCamelCase__            =     time.time()
										UpperCamelCase__            =     0
										completed_steps += 1
					if step % args.save_checkpoint_steps == 0:
										logger.info("""Evaluating and saving model checkpoint""")
										UpperCamelCase__            =     evaluate(args)
										log_metrics(step, {"""loss/eval""": eval_loss, """perplexity""": perplexity})
										accelerator.wait_for_everyone()
										UpperCamelCase__            =     os.path.join(args.save_dir, f'''step_{step}''')
										accelerator.save_state(save_dir)
										if accelerator.is_main_process:
															hf_repo.push_to_hub(commit_message=f'''step {step}''')
										model.train()
					if completed_steps >= args.max_train_steps:
										break

# Evaluate and save the last checkpoint
logger.info("""Evaluating and saving model after training""")
UpperCamelCase__            =     evaluate(args)
log_metrics(step, {"""loss/eval""": eval_loss, """perplexity""": perplexity})
accelerator.wait_for_everyone()
UpperCamelCase__            =     accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
UpperCamelCase__            =     os.path.join(args.save_dir, f'''step_{step}''')
accelerator.save_state(save_dir)
if accelerator.is_main_process:
					hf_repo.push_to_hub(commit_message="""final model""")








import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
       from transformers.models.electra.modeling_flax_electra import (
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForMultipleChoice,
           FlaxElectraForPreTraining,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForSequenceClassification,
           FlaxElectraForTokenClassification,
           FlaxElectraModel,
       )





class 			lowerCAmelCase__(     unittest.TestCase    ):






    '''simple docstring'''





    def __init__(						self					: Any							, __snake_case					: Tuple							, __snake_case					: Union[str, Any]=13							, __snake_case					: Tuple=7							, __snake_case					: str=True							, __snake_case					: List[Any]=True							, __snake_case					: Optional[int]=True							, __snake_case					: List[str]=True							, __snake_case					: Tuple=99							, __snake_case					: int=24							, __snake_case					: int=32							, __snake_case					: Tuple=5							, __snake_case					: int=4							, __snake_case					: Any=37							, __snake_case					: int="gelu"							, __snake_case					: Union[str, Any]=0.1							, __snake_case					: Optional[Any]=0.1							, __snake_case					: str=512							, __snake_case					: List[Any]=16							, __snake_case					: Tuple=2							, __snake_case					: List[str]=0.02							, __snake_case					: Any=4							, ):





       '''simple docstring'''
       UpperCAmelCase_				:   Union[str, Any]							= parent
       UpperCAmelCase_				:   Dict							= batch_size
       UpperCAmelCase_				:   Union[str, Any]							= seq_length
       UpperCAmelCase_				:   Any							= is_training
       UpperCAmelCase_				:   Tuple							= use_attention_mask
       UpperCAmelCase_				:   Dict							= use_token_type_ids
       UpperCAmelCase_				:   str							= use_labels
       UpperCAmelCase_				:   Union[str, Any]							= vocab_size
       UpperCAmelCase_				:   Union[str, Any]							= hidden_size
       UpperCAmelCase_				:   Union[str, Any]							= embedding_size
       UpperCAmelCase_				:   Optional[int]							= num_hidden_layers
       UpperCAmelCase_				:   Tuple							= num_attention_heads
       UpperCAmelCase_				:   List[Any]							= intermediate_size
       UpperCAmelCase_				:   Union[str, Any]							= hidden_act
       UpperCAmelCase_				:   Any							= hidden_dropout_prob
       UpperCAmelCase_				:   List[str]							= attention_probs_dropout_prob
       UpperCAmelCase_				:   List[str]							= max_position_embeddings
       UpperCAmelCase_				:   Dict							= type_vocab_size
       UpperCAmelCase_				:   Tuple							= type_sequence_label_size
       UpperCAmelCase_				:   Dict							= initializer_range
       UpperCAmelCase_				:   Any							= num_choices





    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Any							= ids_tensor([self.batch_size, self.seq_length]							, self.vocab_size						)

       UpperCAmelCase_				:   Any							= None
       if self.use_attention_mask:
          UpperCAmelCase_				:   List[str]							= random_attention_mask([self.batch_size, self.seq_length]						)

       UpperCAmelCase_				:   List[str]							= None
       if self.use_token_type_ids:
          UpperCAmelCase_				:   str							= ids_tensor([self.batch_size, self.seq_length]							, self.type_vocab_size						)

       UpperCAmelCase_				:   int							= ElectraConfig(
           vocab_size=self.vocab_size							, hidden_size=self.hidden_size							, embedding_size=self.embedding_size							, num_hidden_layers=self.num_hidden_layers							, num_attention_heads=self.num_attention_heads							, intermediate_size=self.intermediate_size							, hidden_act=self.hidden_act							, hidden_dropout_prob=self.hidden_dropout_prob							, attention_probs_dropout_prob=self.attention_probs_dropout_prob							, max_position_embeddings=self.max_position_embeddings							, type_vocab_size=self.type_vocab_size							, initializer_range=self.initializer_range							, )

       return config, input_ids, token_type_ids, attention_mask





    def        _lowerCamelCase			(						self					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= self.prepare_config_and_inputs()
       UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_				:   Union[str, Any]							= config_and_inputs
       UpperCAmelCase_				:   Optional[Any]							= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict





@require_flax
class 			lowerCAmelCase__(     snake_case__     ,     unittest.TestCase    ):






    '''simple docstring'''





    A_	:  Any								=			True

    A_	:  int								=			(
        (
            FlaxElectraModel,
            FlaxElectraForCausalLM,
            FlaxElectraForMaskedLM,
            FlaxElectraForPreTraining,
            FlaxElectraForTokenClassification,
            FlaxElectraForQuestionAnswering,
            FlaxElectraForMultipleChoice,
            FlaxElectraForSequenceClassification,
        )
        if is_flax_available()
        else ()
    )





    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= FlaxElectraModelTester(self						)





    @slow
    def        _lowerCamelCase			(						self					: Union[str, Any]						):





       '''simple docstring'''
       for model_class_name in self.all_model_classes:
          if model_class_name == FlaxElectraForMaskedLM:
             UpperCAmelCase_				:   List[str]							= model_class_name.from_pretrained('''google/electra-small-generator'''						)
          else:
             UpperCAmelCase_				:   Optional[Any]							= model_class_name.from_pretrained('''google/electra-small-discriminator'''						)
          UpperCAmelCase_				:   Tuple							= model(np.ones((1, 1)						)						)
          self.assertIsNotNone(__snake_case						)