

"""simple docstring"""






import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device


if is_torch_available():
						import torch

						from transformers import XLMProphetNetForConditionalGeneration, XLMProphetNetTokenizer






@require_torch
class       _UpperCAmelCase     (						unittest.TestCase):
			@slow
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[Any]       =   XLMProphetNetForConditionalGeneration.from_pretrained("""microsoft/xprophetnet-large-wiki100-cased"""       )
				model.to(_A       )

				# encoder-decoder outputs
				_UpperCAmelCase      :							Tuple       =   torch.tensor([[17, 9_62_08, 10_34_71, 2]]       ).to(_A       )
				_UpperCAmelCase      :							Optional[int]       =   torch.tensor(
				    [[2, 2_50, 99_53, 34, 6_94_89, 16_20, 32, 11_84_24, 6_24, 2_10, 1_05, 29_13, 10_32, 3_51]]       ).to(_A       )
				_UpperCAmelCase      :							Optional[Any]       =   model(
				    input_ids=_A ,						attention_mask=_A ,						encoder_outputs=_A ,						decoder_input_ids=_A       )
				_UpperCAmelCase      :							Optional[int]       =   output[0]
				_UpperCAmelCase      :							Optional[Any]       =   torch.Size((1, 14, 25_00_12)       )
				self.assertEqual(output_predited_logis.shape ,						_A       )
				_UpperCAmelCase      :							List[str]       =   torch.tensor(
				    [[[-6.3986, -8.2391, 12.5189], [-6.3289, -8.0864, 12.6211], [-6.2418, -8.0445, 12.7968]]]       ).to(_A       )
				self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3] ,						_A ,						atol=1e-4       )       )

				# encoder outputs
				_UpperCAmelCase      :							Any       =   model.prophetnet.encoder(_A       )[0]
				_UpperCAmelCase      :							int       =   torch.tensor(
				    [[[-1.4260, -0.7628, 0.8453], [-1.4719, -0.1391, 0.7807], [-1.7678, 0.0114, 0.4646]]]       ).to(_A       )
				_UpperCAmelCase      :							int       =   torch.Size((1, 4, 10_24)       )
				self.assertEqual(encoder_outputs.shape ,						_A       )
				self.assertTrue(torch.allclose(encoder_outputs[:, :3, :3] ,						_A ,						atol=1e-4       )       )

				# decoder outputs
				_UpperCAmelCase      :							Optional[int]       =   model.prophetnet.decoder(
				    _A ,						encoder_hidden_states=_A ,						)
				_UpperCAmelCase      :							Optional[int]       =   decoder_outputs[1].view(1 ,						model.config.ngram ,						14 ,						-1       )
				_UpperCAmelCase      :							Tuple       =   model.lm_head(_A       )
				_UpperCAmelCase      :							str       =   predicting_streams_logits[:, 0]
				self.assertTrue(torch.allclose(next_first_stream_logits[:, :3, :3] ,						_A ,						atol=1e-4       )       )
			@slow
			def 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[Any]       =   XLMProphetNetForConditionalGeneration.from_pretrained(
				    """microsoft/xprophetnet-large-wiki100-cased-xglue-ntg"""       )
				model.to(_A       )

				_UpperCAmelCase      :							List[Any]       =   torch.tensor([[17, 9_62_08, 10_34_71, 2]]       ).to(_A       )
				_UpperCAmelCase      :							Optional[int]       =   torch.tensor(
				    [[2, 2_50, 99_53, 34, 6_94_89, 16_20, 32, 11_84_24, 6_24, 2_10, 1_05, 29_13, 10_32, 3_51]]       ).to(_A       )
				_UpperCAmelCase      :							Tuple       =   model(
				    input_ids=_A ,						attention_mask=_A ,						encoder_outputs=_A ,						decoder_input_ids=_A       )
				_UpperCAmelCase      :							int       =   output[0]
				_UpperCAmelCase      :							Optional[int]       =   torch.Size((1, 14, 25_00_12)       )
				self.assertEqual(output_predited_logis.shape ,						_A       )
				# compare the actual values for a slice.
				_UpperCAmelCase      :							List[str]       =   torch.tensor(
				    [[[-9.2253, -9.7173, -6.3529], [-7.6701, -9.0145, -1.9382], [-8.0195, -7.0004, -0.1523]]]       ).to(_A       )

				self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3] ,						_A ,						atol=1e-4       )       )
			@slow
			def 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   XLMProphetNetForConditionalGeneration.from_pretrained(
				    """microsoft/xprophetnet-large-wiki100-cased-xglue-ntg"""       )
				model.to(_A       )
				_UpperCAmelCase      :							Optional[Any]       =   5_12

				_UpperCAmelCase      :							Optional[Any]       =   XLMProphetNetTokenizer.from_pretrained("""microsoft/xprophetnet-large-wiki100-cased-xglue-ntg"""       )

				_UpperCAmelCase      :							Union[str, Any]       =   (
				    """Microsoft Corporation intends to officially end free support for the Windows 7 operating system after"""
				    """ January 14, 2020, according to the official portal of the organization. From that day, users of this"""
				    """ system will not be able to receive security updates, which could make their computers vulnerable to"""
				    """ cyber attacks."""
				)
				_UpperCAmelCase      :							List[str]       =   (
				    """орпорация Microsoft намерена официально прекратить бесплатную поддержку операционной системы Windows 7"""
				    """ после 14 января 2020 года, сообщается на официальном портале организации . С указанного дня пользователи"""
				    """ этой системы не смогут получать обновления безопасности, из-за чего их компьютеры могут стать уязвимыми"""
				    """ к кибератакам."""
				)
				_UpperCAmelCase      :							Optional[Any]       =   (
				    """根据该组织的官方门户网站，微软公司打算在2020年1月14日之后正式终止对Windows 7操作系统的免费支持。从那时起，该系统的用户将无法接收安全更新，这可能会使他们的计算机容易受到网络攻击。"""
				)

				_UpperCAmelCase      :							Optional[Any]       =   tokenizer(
				    [EN_SENTENCE, RU_SENTENCE, ZH_SENTENCE] ,						padding=_A ,						max_length=2_55 ,						return_tensors="""pt"""       ).input_ids
				_UpperCAmelCase      :							List[str]       =   input_ids.to(_A       )

				_UpperCAmelCase      :							Optional[Any]       =   model.generate(
				    _A ,						num_beams=10 ,						length_penalty=1.0 ,						no_repeat_ngram_size=3 ,						early_stopping=_A       )
				_UpperCAmelCase      :							str       =   [tokenizer.decode(_A ,						skip_special_tokens=_A       ) for g in summary_ids]
				_UpperCAmelCase      :							List[str]       =   """Microsoft to end Windows 7 free support after January 14, 2020"""
				_UpperCAmelCase      :							Dict       =   """Microsoft намерена прекратить бесплатную поддержку Windows 7 после 14 января 2020 года"""
				_UpperCAmelCase      :							Optional[int]       =   """微软打算终止对Windows 7操作系统的免费支持"""
				self.assertListEqual(
				    [EXPECTED_TITLE_EN, EXPECTED_TITLE_RU, EXPECTED_TITLE_ZH] ,						_A ,						)

				_UpperCAmelCase      :							List[str]       =   model.generate(
				    _A ,						num_beams=1 ,						length_penalty=1.0 ,						no_repeat_ngram_size=3 ,						early_stopping=_A       )
				_UpperCAmelCase      :							Tuple       =   [
				    tokenizer.convert_ids_to_tokens(_A ,						skip_special_tokens=_A       ) for g in summary_ids_beama
				]
				_UpperCAmelCase      :							List[str]       =   """▁Microsoft ▁to ▁end ▁free ▁support ▁for ▁Windows ▁7""".split(""" """       )
				_UpperCAmelCase      :							Dict       =   """▁Microsoft ▁намерен а ▁прекрати ть ▁бес плат ную ▁поддержку ▁Windows ▁7 ▁после ▁14 ▁января ▁2020 ▁года""".split(
				    """ """       )
				_UpperCAmelCase      :							Union[str, Any]       =   """微软 公司 打算 终止 对 Windows ▁7 操作 系统的 免费 支持""".split(""" """       )
				self.assertListEqual(
				    [EXPECTED_TITLE_EN_BEAM1_TOK, EXPECTED_TITLE_RU_BEAM1_TOK, EXPECTED_TITLE_ZH_BEAM1_TOK] ,						_A ,						)








"""simple docstring"""






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,
    RagRetriever,
    RagSequenceForGeneration,
    RagTokenizer,
)


lowerCamelCase__			:					int							  =					logging.getLogger(__name__)
torch.set_grad_enabled(False)
lowerCamelCase__			:					int							  =					'''cuda''' if torch.cuda.is_available() else '''cpu'''
def  UpperCamelCase   (     _lowerCAmelCase       :  str,					_lowerCAmelCase       :  int=100,					_lowerCAmelCase       :  List[str]=" "	)			->     List[str]:
	_UpperCAmelCase      :							Tuple       =   text.split(_lowerCAmelCase	)
	return [character.join(text[i : i + n]	).strip() for i in range(0,					len(_lowerCAmelCase	),					_lowerCAmelCase	)]
def  UpperCamelCase   (     _lowerCAmelCase       :  dict	)			->     dict:
	_UpperCAmelCase  ,				_UpperCAmelCase      :							List[str]       =   [], []
	for title, text in zip(documents["""title"""],					documents["""text"""]	):
		if text is not None:
			for passage in split_text(_lowerCAmelCase	):
				titles.append(title if title is not None else """"""	)
				texts.append(_lowerCAmelCase	)
	return {"title": titles, "text": texts}
def  UpperCamelCase   (     _lowerCAmelCase       :  dict,					_lowerCAmelCase       :  DPRContextEncoder,					_lowerCAmelCase       :  DPRContextEncoderTokenizerFast	)			->     dict:
	_UpperCAmelCase      :							int       =   ctx_tokenizer(
	    documents["""title"""],					documents["""text"""],					truncation=_lowerCAmelCase,					padding="""longest""",					return_tensors="""pt"""	)["""input_ids"""]
	_UpperCAmelCase      :							Optional[Any]       =   ctx_encoder(input_ids.to(device=_lowerCAmelCase	),					return_dict=_lowerCAmelCase	).pooler_output
	return {"embeddings": embeddings.detach().cpu().numpy()}
def  UpperCamelCase   (     _lowerCAmelCase       :  "RagExampleArguments",					_lowerCAmelCase       :  "ProcessingArguments",					_lowerCAmelCase       :  "IndexHnswArguments",					)			->     int:
	######################################
	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
	_UpperCAmelCase      :							Optional[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
	_UpperCAmelCase      :							List[str]       =   dataset.map(_lowerCAmelCase,					batched=_lowerCAmelCase,					num_proc=processing_args.num_proc	)

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

	# And finally save your dataset
	_UpperCAmelCase      :							Dict       =   os.path.join(rag_example_args.output_dir,					"""my_knowledge_dataset"""	)
	dataset.save_to_disk(_lowerCAmelCase	)
	# 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
	_UpperCAmelCase      :							List[Any]       =   faiss.IndexHNSWFlat(index_hnsw_args.d,					index_hnsw_args.m,					faiss.METRIC_INNER_PRODUCT	)
	dataset.add_faiss_index("""embeddings""",					custom_index=_lowerCAmelCase	)

	# And save the index
	_UpperCAmelCase      :							List[str]       =   os.path.join(rag_example_args.output_dir,					"""my_knowledge_dataset_hnsw_index.faiss"""	)
	dataset.get_index("""embeddings"""	).save(_lowerCAmelCase	)
	# dataset.load_faiss_index("embeddings", index_path)  # to reload the index

	######################################
	logger.info("""Step 3 - Load RAG"""	)
	######################################

	# Easy way to load the model
	_UpperCAmelCase      :							List[Any]       =   RagRetriever.from_pretrained(
	    rag_example_args.rag_model_name,					index_name="""custom""",					indexed_dataset=_lowerCAmelCase	)
	_UpperCAmelCase      :							Dict       =   RagSequenceForGeneration.from_pretrained(rag_example_args.rag_model_name,					retriever=_lowerCAmelCase	)
	_UpperCAmelCase      :							Tuple       =   RagTokenizer.from_pretrained(rag_example_args.rag_model_name	)

	# For distributed fine-tuning you'll need to provide the paths instead, as the dataset and the index are loaded separately.
	# retriever = RagRetriever.from_pretrained(rag_model_name, index_name="custom", passages_path=passages_path, index_path=index_path)

	######################################
	logger.info("""Step 4 - Have fun"""	)
	######################################

	_UpperCAmelCase      :							List[Any]       =   rag_example_args.question or """What does Moses' rod turn into ?"""
	_UpperCAmelCase      :							List[Any]       =   tokenizer.question_encoder(_lowerCAmelCase,					return_tensors="""pt"""	)["""input_ids"""]
	_UpperCAmelCase      :							Dict       =   model.generate(_lowerCAmelCase	)
	_UpperCAmelCase      :							int       =   tokenizer.batch_decode(_lowerCAmelCase,					skip_special_tokens=_lowerCAmelCase	)[0]
	logger.info("""Q: """ + question	)
	logger.info("""A: """ + generated_string	)






@dataclass
class       _UpperCAmelCase     :
			__a	:    str			   =							field(
			    default=str(Path(__a).parent / """test_data""" / """my_knowledge_dataset.csv""")      ,       metadata={"""help""": """Path to a tab-separated csv file with columns 'title' and 'text'"""}      ,       )
			__a	:    Optional[str]			   =							field(
			    default=__a      ,       metadata={"""help""": """Question that is passed as input to RAG. Default is 'What does Moses' rod turn into ?'."""}      ,       )
			__a	:    str			   =							field(
			    default="""facebook/rag-sequence-nq"""      ,       metadata={"""help""": """The RAG model to use. Either 'facebook/rag-sequence-nq' or 'facebook/rag-token-nq'"""}      ,       )
			__a	:    str			   =							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	:    Optional[str]			   =							field(
			    default=__a      ,       metadata={"""help""": """Path to a directory where the dataset passages and the index will be saved"""}      ,       )






@dataclass
class       _UpperCAmelCase     :
			__a	:    Optional[int]			   =							field(
			    default=__a      ,       metadata={
			        """help""": """The number of processes to use to split the documents into passages. Default is single process."""
			    }      ,       )
			__a	:    int			   =							field(
			    default=1_6      ,       metadata={
			        """help""": """The batch size to use when computing the passages embeddings using the DPR context encoder."""
			    }      ,       )






@dataclass
class       _UpperCAmelCase     :
			__a	:    int			   =							field(
			    default=7_6_8      ,       metadata={"""help""": """The dimension of the embeddings to pass to the HNSW Faiss index."""}      ,       )
			__a	:    int			   =							field(
			    default=1_2_8      ,       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)

						lowerCamelCase__			:					Tuple							  =					HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
						lowerCamelCase__  ,					lowerCamelCase__  ,					lowerCamelCase__			:					Dict							  =					parser.parse_args_into_dataclasses()
						with TemporaryDirectory() as tmp_dir:
												lowerCamelCase__			:					List[str]							  =					rag_example_args.output_dir or tmp_dir
												main(rag_example_args, processing_args, index_hnsw_args)





