






import logging
import os
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import List, Optional

import faiss
import torch
from datasets import Features, Sequence, Value, load_dataset

from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast, HfArgumentParser


_a    :   Union[str, Any]        =			logging.getLogger(__name__)
torch.set_grad_enabled(False)
_a    :   Union[str, Any]        =			"""cuda""" if torch.cuda.is_available() else """cpu"""


def 				snake_case__							(							UpperCAmelCase							:  str    ,			UpperCAmelCase							:  Optional[Any]=1_0_0    ,			UpperCAmelCase							:  List[Any]=" "       ):
    lowerCAmelCase__   :List[str]      							=				text.split(UpperCAmelCase       )
    return [character.join(text[i : i + n]       ).strip() for i in range(0    ,			len(UpperCAmelCase       )    ,			UpperCAmelCase       )]


def 				snake_case__							(							UpperCAmelCase							:  dict       ):
    lowerCAmelCase__   :Any      							=				[], []
    for title, text in zip(documents["title"]    ,			documents["text"]       ):
        if text is not None:
            for passage in split_text(UpperCAmelCase       ):
                titles.append(title if title is not None else ""       )
                texts.append(UpperCAmelCase       )
    return {"title": titles, "text": texts}


def 				snake_case__							(							UpperCAmelCase							:  dict    ,			UpperCAmelCase							:  DPRContextEncoder    ,			UpperCAmelCase							:  DPRContextEncoderTokenizerFast       ):
    lowerCAmelCase__   :Optional[int]      							=				ctx_tokenizer(
        documents["title"]    ,			documents["text"]    ,			truncation=UpperCAmelCase    ,			padding="longest"    ,			return_tensors="pt"       )["input_ids"]
    lowerCAmelCase__   :Tuple      							=				ctx_encoder(input_ids.to(device=UpperCAmelCase       )    ,			return_dict=UpperCAmelCase       ).pooler_output
    return {"embeddings": embeddings.detach().cpu().numpy()}


def 				snake_case__							(							UpperCAmelCase							:  "RagExampleArguments"    ,			UpperCAmelCase							:  "ProcessingArguments"    ,			UpperCAmelCase							:  "IndexHnswArguments"    ,			):
    ######################################
    logger.info("Step 1 - Create the dataset"       )
    ######################################

    # The dataset needed for RAG must have three columns:
    # - title (string): title of the document
    # - text (string): text of a passage of the document
    # - embeddings (array of dimension d): DPR representation of the passage
    # Let's say you have documents in tab-separated csv files with columns "title" and "text"
    assert os.path.isfile(rag_example_args.csv_path       ), "Please provide a valid path to a csv file"

    # You can load a Dataset object this way
    lowerCAmelCase__   :int      							=				load_dataset(
        "csv"    ,			data_files=[rag_example_args.csv_path]    ,			split="train"    ,			delimiter="\t"    ,			column_names=["title", "text"]       )

    # More info about loading csv files in the documentation: https://huggingface.co/docs/datasets/loading_datasets.html?highlight=csv#csv-files

    # Then split the documents into passages of 100 words
    lowerCAmelCase__   :Optional[Any]      							=				dataset.map(UpperCAmelCase    ,			batched=UpperCAmelCase    ,			num_proc=processing_args.num_proc       )

    # And compute the embeddings
    lowerCAmelCase__   :List[str]      							=				DPRContextEncoder.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name       ).to(device=UpperCAmelCase       )
    lowerCAmelCase__   :Optional[int]      							=				DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name       )
    lowerCAmelCase__   :List[str]      							=				Features(
        {"text": Value("string"       ), "title": Value("string"       ), "embeddings": Sequence(Value("float32"       )       )}       )  # optional, save as float32 instead of float64 to save space
    lowerCAmelCase__   :Optional[Any]      							=				dataset.map(
        partial(UpperCAmelCase    ,			ctx_encoder=UpperCAmelCase    ,			ctx_tokenizer=UpperCAmelCase       )    ,			batched=UpperCAmelCase    ,			batch_size=processing_args.batch_size    ,			features=UpperCAmelCase    ,			)

    # And finally save your dataset
    lowerCAmelCase__   :int      							=				os.path.join(rag_example_args.output_dir    ,			"my_knowledge_dataset"       )
    dataset.save_to_disk(UpperCAmelCase       )
    # from datasets import load_from_disk
    # dataset = load_from_disk(passages_path)  # to reload the dataset

    ######################################
    logger.info("Step 2 - Index the dataset"       )
    ######################################

    # Let's use the Faiss implementation of HNSW for fast approximate nearest neighbor search
    lowerCAmelCase__   :str      							=				faiss.IndexHNSWFlat(index_hnsw_args.d    ,			index_hnsw_args.m    ,			faiss.METRIC_INNER_PRODUCT       )
    dataset.add_faiss_index("embeddings"    ,			custom_index=UpperCAmelCase       )

    # And save the index
    lowerCAmelCase__   :Optional[Any]      							=				os.path.join(rag_example_args.output_dir    ,			"my_knowledge_dataset_hnsw_index.faiss"       )
    dataset.get_index("embeddings"       ).save(UpperCAmelCase       )
    # dataset.load_faiss_index("embeddings", index_path)  # to reload the index







@dataclass
class 	_UpperCAmelCase  :




      """simple docstring"""



      A          =    field(
          default=str(Path(_A		).parent / '''test_run''' / '''dummy-kb''' / '''my_knowledge_dataset.csv'''		)				,							metadata={'''help''': '''Path to a tab-separated csv file with columns \'title\' and \'text\''''}				,							)
      A          =    field(
          default=_A				,							metadata={'''help''': '''Question that is passed as input to RAG. Default is \'What does Moses\' rod turn into ?\'.'''}				,							)
      A          =    field(
          default='''facebook/rag-sequence-nq'''				,							metadata={'''help''': '''The RAG model to use. Either \'facebook/rag-sequence-nq\' or \'facebook/rag-token-nq\''''}				,							)
      A          =    field(
          default='''facebook/dpr-ctx_encoder-multiset-base'''				,							metadata={
              '''help''': (
                  '''The DPR context encoder model to use. Either \'facebook/dpr-ctx_encoder-single-nq-base\' or'''
                  ''' \'facebook/dpr-ctx_encoder-multiset-base\''''
              )
          }				,							)
      A          =    field(
          default=str(Path(_A		).parent / '''test_run''' / '''dummy-kb'''		)				,							metadata={'''help''': '''Path to a directory where the dataset passages and the index will be saved'''}				,							)







@dataclass
class 	_UpperCAmelCase  :




      """simple docstring"""



      A          =    field(
          default=_A				,							metadata={
              '''help''': '''The number of processes to use to split the documents into passages. Default is single process.'''
          }				,							)
      A          =    field(
          default=16				,							metadata={
              '''help''': '''The batch size to use when computing the passages embeddings using the DPR context encoder.'''
          }				,							)







@dataclass
class 	_UpperCAmelCase  :




      """simple docstring"""



      A          =    field(
          default=7_68				,							metadata={'''help''': '''The dimension of the embeddings to pass to the HNSW Faiss index.'''}				,							)
      A          =    field(
          default=1_28				,							metadata={
              '''help''': (
                  '''The number of bi-directional links created for every new element during the HNSW index construction.'''
              )
          }				,							)


if __name__ == "__main__":
 logging.basicConfig(level=logging.WARNING)
 logger.setLevel(logging.INFO)

 _a    :   Optional[Any]        =			HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
 _a    :   Optional[Any]        =			parser.parse_args_into_dataclasses()
 with TemporaryDirectory() as tmp_dir:
  _a    :   Optional[Any]        =			rag_example_args.output_dir or tmp_dir
  main(rag_example_args, processing_args, index_hnsw_args)













import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import UperNetForSemanticSegmentation
 from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor







class 	_UpperCAmelCase  :




      """simple docstring"""



      def __init__(   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=13    ,	_lowerCAmelCase=32    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=4    ,	_lowerCAmelCase=[10, 20, 30, 40]    ,	_lowerCAmelCase=[2, 2, 3, 2]    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=37    ,	_lowerCAmelCase="gelu"    ,	_lowerCAmelCase=10    ,	_lowerCAmelCase=0.02    ,	_lowerCAmelCase=["stage2", "stage3", "stage4"]    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=None    ,	):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				parent
          lowerCAmelCase__   :Dict      							=				batch_size
          lowerCAmelCase__   :str      							=				image_size
          lowerCAmelCase__   :int      							=				num_channels
          lowerCAmelCase__   :Tuple      							=				num_stages
          lowerCAmelCase__   :Optional[int]      							=				hidden_sizes
          lowerCAmelCase__   :Tuple      							=				depths
          lowerCAmelCase__   :Tuple      							=				is_training
          lowerCAmelCase__   :str      							=				use_labels
          lowerCAmelCase__   :int      							=				intermediate_size
          lowerCAmelCase__   :List[str]      							=				hidden_act
          lowerCAmelCase__   :Dict      							=				type_sequence_label_size
          lowerCAmelCase__   :Any      							=				initializer_range
          lowerCAmelCase__   :List[str]      							=				out_features
          lowerCAmelCase__   :List[Any]      							=				num_labels
          lowerCAmelCase__   :str      							=				scope
          lowerCAmelCase__   :Dict      							=				num_stages
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :int      							=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

          lowerCAmelCase__   :List[str]      							=				None
          if self.use_labels:
              lowerCAmelCase__   :str      							=				ids_tensor([self.batch_size]    ,	self.type_sequence_label_size		)

          lowerCAmelCase__   :Optional[Any]      							=				self.get_config()

          return config, pixel_values, labels
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return ConvNextConfig(
              num_channels=self.num_channels    ,	num_stages=self.num_stages    ,	hidden_sizes=self.hidden_sizes    ,	depths=self.depths    ,	is_training=self.is_training    ,	intermediate_size=self.intermediate_size    ,	hidden_act=self.hidden_act    ,	out_features=self.out_features    ,	)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return UperNetConfig(
              backbone_config=self.get_backbone_config()    ,	hidden_size=512    ,	pool_scales=[1, 2, 3, 6]    ,	use_auxiliary_head=_lowerCAmelCase    ,	auxiliary_loss_weight=0.4    ,	auxiliary_in_channels=40    ,	auxiliary_channels=256    ,	auxiliary_num_convs=1    ,	auxiliary_concat_input=_lowerCAmelCase    ,	loss_ignore_index=255    ,	num_labels=self.num_labels    ,	)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Dict      							=				UperNetForSemanticSegmentation(config=_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          lowerCAmelCase__   :List[Any]      							=				model(_lowerCAmelCase		)
          self.parent.assertEqual(
              result.logits.shape    ,	(self.batch_size, self.num_labels, self.image_size, self.image_size)		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[int]      							=				self.prepare_config_and_inputs()
          (
              (
              lowerCAmelCase__
          )							,(
              lowerCAmelCase__
          )							,(
              lowerCAmelCase__
          )							,
          )   :Dict      							=				config_and_inputs
          lowerCAmelCase__   :Any      							=				{"pixel_values": pixel_values}
          return config, inputs_dict







@require_torch
class 	_UpperCAmelCase  (  _A				,							_A				,							unittest.TestCase		):




      """simple docstring"""



      A          =    (UperNetForSemanticSegmentation,) if is_torch_available() else ()
      A          =    {'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
      A          =    False
      A          =    False
      A          =    False
      A          =    False
      A          =    False
      A          =    False
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				UperNetModelTester(self		)
          lowerCAmelCase__   :Any      							=				ConfigTester(self    ,	config_class=_lowerCAmelCase    ,	has_text_modality=_lowerCAmelCase    ,	hidden_size=37		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.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.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :str      							=				self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              lowerCAmelCase__   :Optional[int]      							=				model_class(_lowerCAmelCase		)
              lowerCAmelCase__   :Tuple      							=				inspect.signature(model.forward		)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              lowerCAmelCase__   :Optional[Any]      							=				[*signature.parameters.keys()]

              lowerCAmelCase__   :Union[str, Any]      							=				["pixel_values"]
              self.assertListEqual(arg_names[:1]    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :int      							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_semantic_segmentation(*_lowerCAmelCase		)
      @unittest.skip(reason="UperNet does not use inputs_embeds"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip(reason="UperNet does not support input and output embeddings"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip(reason="UperNet does not have a base model"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip(reason="UperNet does not have a base model"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @require_torch_multi_gpu
      @unittest.skip(reason="UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      def 			snake_case_      (   self		):



          '''simple docstring'''
          def check_hidden_states_output(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):
              lowerCAmelCase__   :int      							=				model_class(_lowerCAmelCase		)
              model.to(_lowerCAmelCase		)
              model.eval()

              with torch.no_grad():
                  lowerCAmelCase__   :Dict      							=				model(**self._prepare_for_class(_lowerCAmelCase    ,	_lowerCAmelCase		)		)

              lowerCAmelCase__   :int      							=				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

              lowerCAmelCase__   :Tuple      							=				self.model_tester.num_stages
              self.assertEqual(len(_lowerCAmelCase		)    ,	expected_num_stages + 1		)

              # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
              self.assertListEqual(
                  list(hidden_states[0].shape[-2:]		)    ,	[self.model_tester.image_size // 4, self.model_tester.image_size // 4]    ,	)

          lowerCAmelCase__							,lowerCAmelCase__   :Any      							=				self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              lowerCAmelCase__   :str      							=				True
              check_hidden_states_output(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              lowerCAmelCase__   :str      							=				True

              check_hidden_states_output(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :Dict      							=				self.model_tester.prepare_config_and_inputs_for_common()

          lowerCAmelCase__   :str      							=				_config_zero_init(_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				_config_zero_init(configs_no_init.backbone_config		)
          for model_class in self.all_model_classes:
              lowerCAmelCase__   :List[Any]      							=				model_class(config=_lowerCAmelCase		)
              for name, param in model.named_parameters():
                  if param.requires_grad:
                      self.assertIn(
                          ((param.data.mean() * 1e9).round() / 1e9).item()    ,	[0.0, 1.0]    ,	msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''    ,	)
      @unittest.skip(reason="UperNet does not have tied weights"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass



      @slow
      def 			snake_case_      (   self		):



          '''simple docstring'''
          for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              lowerCAmelCase__   :List[str]      							=				UperNetForSemanticSegmentation.from_pretrained(_lowerCAmelCase		)
              self.assertIsNotNone(_lowerCAmelCase		)





def 				snake_case__							(							):
    lowerCAmelCase__   :int      							=				hf_hub_download(
        repo_id="hf-internal-testing/fixtures_ade20k"    ,			repo_type="dataset"    ,			filename="ADE_val_00000001.jpg"       )
    lowerCAmelCase__   :Union[str, Any]      							=				Image.open(UpperCAmelCase       ).convert("RGB"       )
    return image







@require_torch
@require_vision
@slow
class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"		)
          lowerCAmelCase__   :Dict      							=				UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"		).to(_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				prepare_img()
          lowerCAmelCase__   :str      							=				processor(images=_lowerCAmelCase    ,	return_tensors="pt"		).to(_lowerCAmelCase		)

          with torch.no_grad():
              lowerCAmelCase__   :int      							=				model(**_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				torch.Size((1, model.config.num_labels, 512, 512)		)
          self.assertEqual(outputs.logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :Tuple      							=				torch.tensor(
              [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]		).to(_lowerCAmelCase		)
          self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]    ,	_lowerCAmelCase    ,	atol=1e-4		)		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"		)
          lowerCAmelCase__   :Optional[int]      							=				UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"		).to(_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				prepare_img()
          lowerCAmelCase__   :Tuple      							=				processor(images=_lowerCAmelCase    ,	return_tensors="pt"		).to(_lowerCAmelCase		)

          with torch.no_grad():
              lowerCAmelCase__   :Union[str, Any]      							=				model(**_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				torch.Size((1, model.config.num_labels, 512, 512)		)
          self.assertEqual(outputs.logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				torch.tensor(
              [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]		).to(_lowerCAmelCase		)
          self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]    ,	_lowerCAmelCase    ,	atol=1e-4		)		)



