code: 
  
  
  
  
  """simple docstring"""
  
  
  
  import json
  import os
  import shutil
  import sys
  import tempfile
  import unittest
  from unittest import TestCase
  from unittest.mock import patch
  
  import faiss
  import numpy as np
  from datasets import Dataset
  
  from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
  from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
  from transformers.integrations import is_ray_available
  from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
  from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
  from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
  from transformers.testing_utils import require_ray
  
  
  sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip
  
  if is_torch_available():
        from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
  else:
        SCREAMING_SNAKE_CASE       = None
  
  if is_ray_available():
        import ray  # noqa: E402 # isort:skip
        from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
  else:
        SCREAMING_SNAKE_CASE       = None
        SCREAMING_SNAKE_CASE       = None
        SCREAMING_SNAKE_CASE       = None
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Optional[Any]:
   if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
    A__ =       unittest.skip("test requires Datasets, Faiss, psutil"		)(lowercase_		)
   return test_case
  
  
  
  
  
  
  @require_distributed_retrieval
  class 			UpperCAmelCase_						(		A_							):
  
  
  
  
      def 							__magic_name__							(					self		: Dict   )   ->							Optional[Any]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       tempfile.mkdtemp()
       A__ =       8
  
       # DPR tok
       A__ =       [
           "[UNK]",
           "[CLS]",
           "[SEP]",
           "[PAD]",
           "[MASK]",
           "want",
           "##want",
           "##ed",
           "wa",
           "un",
           "runn",
           "##ing",
           ",",
           "low",
           "lowest",
       ]
       A__ =       os.path.join(self.tmpdirname			,				"dpr_tokenizer"   )
       os.makedirs(snake_case_			,				exist_ok=snake_case_   )
       A__ =       os.path.join(snake_case_			,				DPR_VOCAB_FILES_NAMES["vocab_file"]   )
       with open(self.vocab_file			,				"w"			,				encoding="utf-8"   ) as vocab_writer:
        vocab_writer.write("".join([x + "\n" for x in vocab_tokens]   )   )
  
       # BART tok
       A__ =       [
           "l",
           "o",
           "w",
           "e",
           "r",
           "s",
           "t",
           "i",
           "d",
           "n",
           "\u0120",
           "\u0120l",
           "\u0120n",
           "\u0120lo",
           "\u0120low",
           "er",
           "\u0120lowest",
           "\u0120newer",
           "\u0120wider",
           "<unk>",
       ]
       A__ =       dict(zip(snake_case_			,				range(len(snake_case_   )   )   )   )
       A__ =       ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
       A__ =       {"unk_token": "<unk>"}
  
       A__ =       os.path.join(self.tmpdirname			,				"bart_tokenizer"   )
       os.makedirs(snake_case_			,				exist_ok=snake_case_   )
       A__ =       os.path.join(snake_case_			,				BART_VOCAB_FILES_NAMES["vocab_file"]   )
       A__ =       os.path.join(snake_case_			,				BART_VOCAB_FILES_NAMES["merges_file"]   )
       with open(self.vocab_file			,				"w"			,				encoding="utf-8"   ) as fp:
        fp.write(json.dumps(snake_case_   ) + "\n"   )
       with open(self.merges_file			,				"w"			,				encoding="utf-8"   ) as fp:
        fp.write("\n".join(snake_case_   )   )
  
  
  
  
      def 							__magic_name__							(					self		: Dict   )   ->							DPRQuestionEncoderTokenizer:
  
  
  
  
       '''simple docstring'''
  
       return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname			,				"dpr_tokenizer"   )   )
  
  
  
  
      def 							__magic_name__							(					self		: Union[str, Any]   )   ->							BartTokenizer:
  
  
  
  
       '''simple docstring'''
  
       return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname			,				"bart_tokenizer"   )   )
  
  
  
  
      def 							__magic_name__							(					self		: Dict   )   ->							Tuple:
  
  
  
  
       '''simple docstring'''
  
       shutil.rmtree(self.tmpdirname   )
  
  
  
  
      def 							__magic_name__							(					self		: List[Any]   )   ->							Tuple:
  
  
  
  
       '''simple docstring'''
  
       A__ =       Dataset.from_dict(
           {
               "id": ["0", "1"],
               "text": ["foo", "bar"],
               "title": ["Foo", "Bar"],
               "embeddings": [np.ones(self.retrieval_vector_size   ), 2 * np.ones(self.retrieval_vector_size   )],
           }   )
       dataset.add_faiss_index("embeddings"			,				string_factory="Flat"			,				metric_type=faiss.METRIC_INNER_PRODUCT   )
       return dataset
  
  
  
  
      def 							__magic_name__							(					self		: Any			,				snake_case_		: bool			,				snake_case_		: List[Any]=12_345   )   ->							RagPyTorchDistributedRetriever:
  
  
  
  
       '''simple docstring'''
  
       A__ =       self.get_dummy_dataset()
       A__ =       RagConfig(
           retrieval_vector_size=self.retrieval_vector_size			,				question_encoder=DPRConfig().to_dict()			,				generator=BartConfig().to_dict()			,				)
       with patch("transformers.models.rag.retrieval_rag.load_dataset"   ) as mock_load_dataset:
        A__ =       dataset
        A__ =       RagPyTorchDistributedRetriever(
            snake_case_			,				question_encoder_tokenizer=self.get_dpr_tokenizer()			,				generator_tokenizer=self.get_bart_tokenizer()			,				)
        if init_retrieval:
         retriever.init_retrieval(snake_case_   )
       return retriever
  
  
  
  
      def 							__magic_name__							(					self		: List[str]			,				snake_case_		: bool   )   ->							RagRayDistributedRetriever:
  
  
  
  
       '''simple docstring'''
  
       ray.init(local_mode=snake_case_   )
       A__ =       RagConfig(
           retrieval_vector_size=self.retrieval_vector_size			,				question_encoder=DPRConfig().to_dict()			,				generator=BartConfig().to_dict()			,				)
       A__ =       ray.remote(snake_case_   )
       A__ =       [remote_cls.remote() for _ in range(1   )]
       with patch("transformers.models.rag.retrieval_rag.load_dataset"   ) as mock_load_dataset:
        A__ =       self.get_dummy_dataset()
        A__ =       RagRayDistributedRetriever(
            snake_case_			,				question_encoder_tokenizer=self.get_dpr_tokenizer()			,				generator_tokenizer=self.get_bart_tokenizer()			,				retrieval_workers=snake_case_			,				)
        if init_retrieval:
         retriever.init_retrieval()
       return retriever
  
  
  
  
      def 							__magic_name__							(					self		: Union[str, Any]			,				snake_case_		: bool			,				snake_case_		: bool			,				snake_case_		: Any=12_345   )   ->							Tuple:
  
  
  
  
       '''simple docstring'''
  
       A__ =       self.get_dummy_dataset()
       A__ =       RagConfig(
           retrieval_vector_size=self.retrieval_vector_size			,				question_encoder=DPRConfig().to_dict()			,				generator=BartConfig().to_dict()			,				index_name="custom"			,				)
       if from_disk:
        A__ =       os.path.join(self.tmpdirname			,				"dataset"   )
        A__ =       os.path.join(self.tmpdirname			,				"index.faiss"   )
        dataset.get_index("embeddings"   ).save(os.path.join(self.tmpdirname			,				"index.faiss"   )   )
        dataset.drop_index("embeddings"   )
        dataset.save_to_disk(os.path.join(self.tmpdirname			,				"dataset"   )   )
        del dataset
        A__ =       RagPyTorchDistributedRetriever(
            snake_case_			,				question_encoder_tokenizer=self.get_dpr_tokenizer()			,				generator_tokenizer=self.get_bart_tokenizer()			,				)
       else:
        A__ =       RagPyTorchDistributedRetriever(
            snake_case_			,				question_encoder_tokenizer=self.get_dpr_tokenizer()			,				generator_tokenizer=self.get_bart_tokenizer()			,				index=CustomHFIndex(config.retrieval_vector_size			,				snake_case_   )			,				)
       if init_retrieval:
        retriever.init_retrieval(snake_case_   )
       return retriever
  
  
  
  
      def 							__magic_name__							(					self		: Any			,				snake_case_		: bool			,				snake_case_		: bool   )   ->							Union[str, Any]:
  
  
  
  
       '''simple docstring'''
  
       ray.init(local_mode=snake_case_   )
       A__ =       self.get_dummy_dataset()
       A__ =       RagConfig(
           retrieval_vector_size=self.retrieval_vector_size			,				question_encoder=DPRConfig().to_dict()			,				generator=BartConfig().to_dict()			,				index_name="custom"			,				)
       A__ =       ray.remote(snake_case_   )
       A__ =       [remote_cls.remote() for _ in range(1   )]
       if from_disk:
        A__ =       os.path.join(self.tmpdirname			,				"dataset"   )
        A__ =       os.path.join(self.tmpdirname			,				"index.faiss"   )
        dataset.get_index("embeddings"   ).save(os.path.join(self.tmpdirname			,				"index.faiss"   )   )
        dataset.drop_index("embeddings"   )
        dataset.save_to_disk(os.path.join(self.tmpdirname			,				"dataset"   )   )
        del dataset
        A__ =       RagRayDistributedRetriever(
            snake_case_			,				question_encoder_tokenizer=self.get_dpr_tokenizer()			,				generator_tokenizer=self.get_bart_tokenizer()			,				retrieval_workers=snake_case_			,				index=CustomHFIndex.load_from_disk(
                vector_size=config.retrieval_vector_size			,				dataset_path=config.passages_path			,				index_path=config.index_path			,				)			,				)
       else:
        A__ =       RagRayDistributedRetriever(
            snake_case_			,				question_encoder_tokenizer=self.get_dpr_tokenizer()			,				generator_tokenizer=self.get_bart_tokenizer()			,				retrieval_workers=snake_case_			,				index=CustomHFIndex(config.retrieval_vector_size			,				snake_case_   )			,				)
       if init_retrieval:
        retriever.init_retrieval()
       return retriever
  
  
  
  
      def 							__magic_name__							(					self		: Any			,				snake_case_		: RagRetriever			,				snake_case_		: np.array			,				snake_case_		: int   )   ->							None:
  
  
  
  
       '''simple docstring'''
  
       A__,   A__,   A__ =       retriever.retrieve(snake_case_			,				n_docs=snake_case_   )
       self.assertEqual(retrieved_doc_embeds.shape			,				(2, n_docs, self.retrieval_vector_size)   )
       self.assertEqual(len(snake_case_   )			,				2   )
       self.assertEqual(sorted(doc_dicts[0]   )			,				["embeddings", "id", "text", "title"]   )
       self.assertEqual(len(doc_dicts[0]["id"]   )			,				snake_case_   )
       self.assertEqual(doc_dicts[0]["id"][0]			,				"1"   )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]			,				"0"   )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()			,				[[1], [0]]   )
  
  
  
  
      def 							__magic_name__							(					self		: str   )   ->							Optional[Any]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       1
       A__ =       np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]			,				dtype=np.floataa   )
  
       self.distributed_retriever_check(
           self.get_dummy_pytorch_distributed_retriever(init_retrieval=snake_case_   )			,				snake_case_			,				snake_case_   )
  
  
  
  
      def 							__magic_name__							(					self		: Dict   )   ->							Any:
  
  
  
  
       '''simple docstring'''
  
       A__ =       1
       A__ =       np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]			,				dtype=np.floataa   )
  
       self.distributed_retriever_check(
           self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=snake_case_			,				from_disk=snake_case_   )			,				snake_case_			,				snake_case_			,				)
  
  
  
  
      def 							__magic_name__							(					self		: Optional[Any]   )   ->							Optional[int]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       1
       A__ =       np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]			,				dtype=np.floataa   )
  
       self.distributed_retriever_check(
           self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=snake_case_			,				from_disk=snake_case_   )			,				snake_case_			,				snake_case_			,				)
  
  
  
  
      @require_ray
      def 							__magic_name__							(					self		: List[str]   )   ->							List[Any]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       1
       A__ =       np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]			,				dtype=np.floataa   )
  
       self.distributed_retriever_check(
           self.get_dummy_ray_distributed_retriever(init_retrieval=snake_case_   )			,				snake_case_			,				snake_case_   )
       ray.shutdown()
  
  
  
  
      @require_ray
      def 							__magic_name__							(					self		: List[Any]   )   ->							Optional[Any]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       1
       A__ =       np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]			,				dtype=np.floataa   )
       with self.assertRaises(snake_case_   ):
        self.distributed_retriever_check(
            self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=snake_case_			,				from_disk=snake_case_   )			,				snake_case_			,				snake_case_			,				)
       ray.shutdown()
  
  
  
  
      @require_ray
      def 							__magic_name__							(					self		: Any   )   ->							Optional[int]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       1
       A__ =       np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]			,				dtype=np.floataa   )
  
       self.distributed_retriever_check(
           self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=snake_case_			,				from_disk=snake_case_   )			,				snake_case_			,				snake_case_   )
       ray.shutdown()
  
  
  
  
  
  
style_context: 
  
  
  
  
  """simple docstring"""
  
  
  
  import logging
  from pathlib import Path
  
  import numpy as np
  import pytorch_lightning as pl
  import torch
  from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
  from pytorch_lightning.utilities import rank_zero_only
  
  from utils import save_json
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	List[Any]:
   A__ =       filter(lambda lowercase_		: p.requires_grad    ,				model.parameters()		)
   A__ =       sum([np.prod(p.size()		) for p in model_parameters]		)
   return params
  
  
  SCREAMING_SNAKE_CASE       = logging.getLogger(__name__)
  
  
  
  
  class 			UpperCAmelCase_						(		pl.Callback							):
  
  
  
  
      def 							__magic_name__							(					self		: List[str]			,				snake_case_		: Optional[Any]			,				snake_case_		: Tuple   )   ->							Optional[int]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       {F"""lr_group_{i}""": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups   )}
       pl_module.logger.log_metrics(snake_case_   )
  
  
  
  
      @rank_zero_only
      def 							__magic_name__							(					self		: Dict			,				snake_case_		: pl.Trainer			,				snake_case_		: pl.LightningModule			,				snake_case_		: str			,				snake_case_		: List[str]=True   )   ->							None:
  
  
  
  
       '''simple docstring'''
  
       logger.info(F"""***** {type_path} results at step {trainer.global_step:05d} *****"""   )
       A__ =       trainer.callback_metrics
       trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]}   )
       # Log results
       A__ =       Path(pl_module.hparams.output_dir   )
       if type_path == "test":
        A__ =       od / "test_results.txt"
        A__ =       od / "test_generations.txt"
       else:
        # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
        # If people want this it will be easy enough to add back.
        A__ =       od / F"""{type_path}_results/{trainer.global_step:05d}.txt"""
        A__ =       od / F"""{type_path}_generations/{trainer.global_step:05d}.txt"""
        results_file.parent.mkdir(exist_ok=snake_case_   )
        generations_file.parent.mkdir(exist_ok=snake_case_   )
       with open(snake_case_			,				"a+"   ) as writer:
        for key in sorted(snake_case_   ):
         if key in ["log", "progress_bar", "preds"]:
          continue
         A__ =       metrics[key]
         if isinstance(snake_case_			,				torch.Tensor   ):
          A__ =       val.item()
         A__ =       F"""{key}: {val:.6f}\n"""
         writer.write(snake_case_   )
  
       if not save_generations:
        return
  
       if "preds" in metrics:
        A__ =       "\n".join(metrics["preds"]   )
        generations_file.open("w+"   ).write(snake_case_   )
  
  
  
  
      @rank_zero_only
      def 							__magic_name__							(					self		: str			,				snake_case_		: Tuple			,				snake_case_		: Dict   )   ->							Any:
  
  
  
  
       '''simple docstring'''
  
       try:
        A__ =       pl_module.model.model.num_parameters()
       except AttributeError:
        A__ =       pl_module.model.num_parameters()
  
       A__ =       count_trainable_parameters(snake_case_   )
       # mp stands for million parameters
       trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6}   )
  
  
  
  
      @rank_zero_only
      def 							__magic_name__							(					self		: List[Any]			,				snake_case_		: pl.Trainer			,				snake_case_		: pl.LightningModule   )   ->							Union[str, Any]:
  
  
  
  
       '''simple docstring'''
  
       save_json(pl_module.metrics			,				pl_module.metrics_save_path   )
       return self._write_logs(snake_case_			,				snake_case_			,				"test"   )
  
  
  
  
      @rank_zero_only
      def 							__magic_name__							(					self		: Optional[int]			,				snake_case_		: pl.Trainer			,				snake_case_		: int   )   ->							Optional[Any]:
  
  
  
  
       '''simple docstring'''
  
       save_json(pl_module.metrics			,				pl_module.metrics_save_path   )
       # Uncommenting this will save val generations
       # return self._write_logs(trainer, pl_module, "valid")
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_    ,				lowercase_=1    ,				lowercase_=False		)		->	Dict:
   if metric == "rouge2":
    A__ =       "{val_avg_rouge2:.4f}-{step_count}"
   elif metric == "bleu":
    A__ =       "{val_avg_bleu:.4f}-{step_count}"
   elif metric == "loss":
    A__ =       "{val_avg_loss:.4f}-{step_count}"
   else:
    raise NotImplementedError(
        f"""seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"""
        " this function."		)
  
   A__ =       ModelCheckpoint(
       dirpath=lowercase_    ,				filename=lowercase_    ,				monitor=f"""val_{metric}"""    ,				mode="min" if "loss" in metric else "max"    ,				save_top_k=lowercase_    ,				)
   return checkpoint_callback
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_		)		->	Union[str, Any]:
   return EarlyStopping(
       monitor=f"""val_{metric}"""    ,				mode="min" if "loss" in metric else "max"    ,				patience=lowercase_    ,				verbose=lowercase_    ,				)
  
  
  
  
  
  