




'''simple docstring'''




import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      Any=2       ,						SCREAMING_SNAKE_CASE__			:      List[str]=True       ,						SCREAMING_SNAKE_CASE__			:      Optional[int]=False       ):
 '''simple docstring'''





 UpperCAmelCase__										=    Accelerator(step_scheduler_with_optimizer=SCREAMING_SNAKE_CASE__       ,						split_batches=SCREAMING_SNAKE_CASE__       )
 UpperCAmelCase__										=    torch.nn.Linear(2       ,						4       )
 UpperCAmelCase__										=    torch.optim.AdamW(model.parameters()       ,						lr=1.0       )
 UpperCAmelCase__										=    torch.optim.lr_scheduler.OneCycleLR(SCREAMING_SNAKE_CASE__       ,						max_lr=0.01       ,						steps_per_epoch=2       ,						epochs=1       )
 UpperCAmelCase__							,			UpperCAmelCase__							,			UpperCAmelCase__										=    accelerator.prepare(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       )

 # Optimizer has stepped
 scheduler.step()
 if step_scheduler_with_optimizer or (num_processes == 1):
  assert (
      scheduler.scheduler.last_epoch == num_processes
  ), F'''Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})'''
 else:
  assert (
      scheduler.scheduler.last_epoch != num_processes
  ), F'''Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})'''

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      str=2       ,						SCREAMING_SNAKE_CASE__			:      List[str]=True       ,						SCREAMING_SNAKE_CASE__			:      Any=False       ):
 '''simple docstring'''





 UpperCAmelCase__										=    Accelerator(step_scheduler_with_optimizer=SCREAMING_SNAKE_CASE__       ,						split_batches=SCREAMING_SNAKE_CASE__       )
 UpperCAmelCase__										=    torch.nn.Linear(2       ,						4       )
 UpperCAmelCase__										=    torch.optim.AdamW(model.parameters()       ,						lr=1.0       )
 UpperCAmelCase__										=    torch.optim.lr_scheduler.LambdaLR(SCREAMING_SNAKE_CASE__       ,						lr_lambda=lambda SCREAMING_SNAKE_CASE__       : 1 - n / 10       )
 UpperCAmelCase__							,			UpperCAmelCase__							,			UpperCAmelCase__										=    accelerator.prepare(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       )

 # Optimizer has stepped
 UpperCAmelCase__										=    False
 scheduler.step()
 UpperCAmelCase__										=    1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), F'''Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'''

 # Optimizer has not stepped
 UpperCAmelCase__										=    True
 scheduler.step()
 if not step_scheduler_with_optimizer:
  UpperCAmelCase__										=    1 - 2 / 10
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), F'''Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'''

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      str = 2       ):
 '''simple docstring'''





 from transformers import get_linear_schedule_with_warmup

 UpperCAmelCase__										=    [1, 2, 4]
 for num_steps in steps:
  UpperCAmelCase__										=    GradientAccumulationPlugin(num_steps=SCREAMING_SNAKE_CASE__       ,						adjust_scheduler=num_steps > 1       )
  UpperCAmelCase__										=    Accelerator(gradient_accumulation_plugin=SCREAMING_SNAKE_CASE__       )
  UpperCAmelCase__										=    torch.nn.Linear(2       ,						4       )
  UpperCAmelCase__										=    torch.optim.AdamW(model.parameters()       ,						lr=10.0       )
  UpperCAmelCase__										=    get_linear_schedule_with_warmup(optimizer=SCREAMING_SNAKE_CASE__       ,						num_warmup_steps=0       ,						num_training_steps=20       )

  UpperCAmelCase__							,			UpperCAmelCase__							,			UpperCAmelCase__										=    accelerator.prepare(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       )

  for i in range(10 * num_steps       ):
   with accelerator.accumulate(SCREAMING_SNAKE_CASE__       ):
    optimizer.step()
    scheduler.step()

   if i == (10 * num_steps - 2):
    assert (
        scheduler.get_last_lr()[0] != 0
    ), F'''Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}'''
  assert (
      scheduler.get_last_lr()[0] == 0
  ), F'''Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}'''
  GradientState._reset_state()







@require_cpu
class 						lowerCAmelCase_   ( unittest.TestCase   ):







  '''simple docstring'''







  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1 )						,		num_processes=1 )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1						,		split_batches=__SCREAMING_SNAKE_CASE )						,		num_processes=1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1 )						,		num_processes=1 )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1						,		split_batches=__SCREAMING_SNAKE_CASE )						,		num_processes=1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict ):






   """simple docstring"""






   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1						,		step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE )						,		num_processes=1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple ):






   """simple docstring"""






   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1						,		step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE )						,		num_processes=1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE )
   debug_launcher(__SCREAMING_SNAKE_CASE )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1						,		split_batches=__SCREAMING_SNAKE_CASE )						,		num_processes=1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE )
   debug_launcher(__SCREAMING_SNAKE_CASE )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1						,		split_batches=__SCREAMING_SNAKE_CASE )						,		num_processes=1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Union[str, Any] ):






   """simple docstring"""






   AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE ) )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






   AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE ) )



  @require_huggingface_suite
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple ):






   """simple docstring"""






   AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE )
   debug_launcher(partial(__SCREAMING_SNAKE_CASE						,		num_processes=1 ) )
   debug_launcher(__SCREAMING_SNAKE_CASE )






'''simple docstring'''




import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


UpperCAmelCase_       				=      logging.getLogger(__name__)

class 						lowerCAmelCase_   :







  '''simple docstring'''







  def __init__(						self    :	int ):






   """simple docstring"""






   UpperCAmelCase__										=    False


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any						,		_UpperCAmelCase    :	Optional[Any]						,		_UpperCAmelCase    :	str						,		_UpperCAmelCase    :	Tuple						,		_UpperCAmelCase    :	Optional[Any] ):






   """simple docstring"""






   if not self.initialized:
    UpperCAmelCase__										=    RagRetriever(
        _UpperCAmelCase						,		question_encoder_tokenizer=_UpperCAmelCase						,		generator_tokenizer=_UpperCAmelCase						,		index=_UpperCAmelCase						,		init_retrieval=_UpperCAmelCase						,		)
    UpperCAmelCase__										=    True


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Union[str, Any] ):






   """simple docstring"""






   self.retriever.index.init_index()


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Union[str, Any] ):






   """simple docstring"""






   del self.retriever
   UpperCAmelCase__										=    False



  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any						,		_UpperCAmelCase    :	Union[str, Any]						,		_UpperCAmelCase    :	Optional[int] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.retriever._main_retrieve(_UpperCAmelCase						,		_UpperCAmelCase )
   UpperCAmelCase__										=    self.retriever.index.get_doc_dicts(_UpperCAmelCase )
   return doc_ids, retrieved_doc_embeds, doc_dicts

class 						lowerCAmelCase_   ( lowerCamelCase_   ):







  '''simple docstring'''







  def __init__(						self    :	Optional[Any]						,		_UpperCAmelCase    :	List[Any]						,		_UpperCAmelCase    :	Dict						,		_UpperCAmelCase    :	str						,		_UpperCAmelCase    :	Any						,		_UpperCAmelCase    :	List[Any]=None ):






   """simple docstring"""






   if index is not None and index.is_initialized() and len(_UpperCAmelCase ) > 0:
    raise ValueError(
        """When using Ray for distributed fine-tuning, """
        """you'll need to provide the paths instead, """
        """as the dataset and the index are loaded """
        """separately. More info in examples/rag/use_own_knowledge_dataset.py """ )

   super().__init__(
       _UpperCAmelCase						,		question_encoder_tokenizer=_UpperCAmelCase						,		generator_tokenizer=_UpperCAmelCase						,		index=_UpperCAmelCase						,		init_retrieval=_UpperCAmelCase						,		)

   UpperCAmelCase__										=    retrieval_workers
   UpperCAmelCase__										=    question_encoder_tokenizer
   UpperCAmelCase__										=    generator_tokenizer
   if len(self.retrieval_workers ) > 0:
    ray.get(
        [
            worker.create_rag_retriever.remote(_UpperCAmelCase						,		_UpperCAmelCase						,		_UpperCAmelCase						,		_UpperCAmelCase )
            for worker in self.retrieval_workers
        ] )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   logger.info("""initializing retrieval""" )

   if len(self.retrieval_workers ) > 0:
    ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers] )
   else:
    # Non-distributed training. Load index into this same process.
    self.index.init_index()


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple						,		_UpperCAmelCase    :	str						,		_UpperCAmelCase    :	Dict ):






   """simple docstring"""






   if len(self.retrieval_workers ) > 0:
    # Select a random retrieval actor.
    UpperCAmelCase__										=    self.retrieval_workers[random.randint(0						,		len(self.retrieval_workers ) - 1 )]
    UpperCAmelCase__							,			UpperCAmelCase__							,			UpperCAmelCase__										=    ray.get(
        random_worker.retrieve.remote(_UpperCAmelCase						,		_UpperCAmelCase ) )
   else:
    UpperCAmelCase__							,			UpperCAmelCase__										=    self._main_retrieve(_UpperCAmelCase						,		_UpperCAmelCase )
    UpperCAmelCase__										=    self.index.get_doc_dicts(_UpperCAmelCase )
   return retrieved_doc_embeds, doc_ids, doc_dicts


  @classmethod
  def 							SCREAMING_SNAKE_CASE__			(						cls    :	str						,		_UpperCAmelCase    :	Optional[Any]						,		_UpperCAmelCase    :	str=None						,		**_UpperCAmelCase    :	Tuple ):






   """simple docstring"""






   return super(_UpperCAmelCase						,		cls ).get_tokenizers(_UpperCAmelCase						,		_UpperCAmelCase						,		**_UpperCAmelCase )


  @classmethod
  def 							SCREAMING_SNAKE_CASE__			(						cls    :	int						,		_UpperCAmelCase    :	Optional[int]						,		_UpperCAmelCase    :	Optional[int]						,		_UpperCAmelCase    :	Union[str, Any]=None						,		**_UpperCAmelCase    :	int ):






   """simple docstring"""






   UpperCAmelCase__										=    kwargs.pop("""config"""						,		_UpperCAmelCase ) or RagConfig.from_pretrained(_UpperCAmelCase						,		**_UpperCAmelCase )
   UpperCAmelCase__										=    RagTokenizer.from_pretrained(_UpperCAmelCase						,		config=_UpperCAmelCase )
   UpperCAmelCase__										=    rag_tokenizer.question_encoder
   UpperCAmelCase__										=    rag_tokenizer.generator

   if indexed_dataset is not None:
    UpperCAmelCase__										=    """custom"""
    UpperCAmelCase__										=    CustomHFIndex(config.retrieval_vector_size						,		_UpperCAmelCase )
   else:
    UpperCAmelCase__										=    cls._build_index(_UpperCAmelCase )

   return cls(
       _UpperCAmelCase						,		question_encoder_tokenizer=_UpperCAmelCase						,		generator_tokenizer=_UpperCAmelCase						,		retrieval_workers=_UpperCAmelCase						,		index=_UpperCAmelCase						,		)



  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict ):






   """simple docstring"""






   logger.info("""re-loading the new dataset with embeddings""" )
   # access from the training loop

   ray.get([worker.clear_object.remote() for worker in self.retrieval_workers] )

   # build the index object again
   UpperCAmelCase__										=    self._build_index(self.config )

   ray.get(
       [
           worker.create_rag_retriever.remote(
               self.config						,		self.question_encoder_tokenizer						,		self.generator_tokenizer						,		_UpperCAmelCase )
           for worker in self.retrieval_workers
       ] )
