


"""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  =	logging.getLogger(__name__)
torch.set_grad_enabled(False)
_lowerCamelCase  =	"cuda" if torch.cuda.is_available() else "cpu"


def __lowercase  (   lowerCamelCase_	:   str    ,       lowerCamelCase_	:   Optional[Any]=100    ,       lowerCamelCase_	:   List[Any]=" "	):
				SCREAMING_SNAKE_CASE__  				=					text.split(lowerCamelCase_	)
				return [character.join(text[i : i + n]	).strip() for i in range(0    ,       len(lowerCamelCase_	)    ,       lowerCamelCase_	)]


def __lowercase  (   lowerCamelCase_	:   dict	):
				SCREAMING_SNAKE_CASE__  				=					[], []
				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 __lowercase  (   lowerCamelCase_	:   dict    ,       lowerCamelCase_	:   DPRContextEncoder    ,       lowerCamelCase_	:   DPRContextEncoderTokenizerFast	):
				SCREAMING_SNAKE_CASE__  				=					ctx_tokenizer(
				    documents["title"]    ,       documents["text"]    ,       truncation=lowerCamelCase_    ,       padding="longest"    ,       return_tensors="pt"	)['''input_ids''']
				SCREAMING_SNAKE_CASE__  				=					ctx_encoder(input_ids.to(device=lowerCamelCase_	)    ,       return_dict=lowerCamelCase_	).pooler_output
				return {"embeddings": embeddings.detach().cpu().numpy()}


def __lowercase  (   lowerCamelCase_	:   "RagExampleArguments"    ,       lowerCamelCase_	:   "ProcessingArguments"    ,       lowerCamelCase_	:   "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
				SCREAMING_SNAKE_CASE__  				=					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
				SCREAMING_SNAKE_CASE__  				=					dataset.map(lowerCamelCase_    ,       batched=lowerCamelCase_    ,       num_proc=processing_args.num_proc	)

				# And compute the embeddings
				SCREAMING_SNAKE_CASE__  				=					DPRContextEncoder.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name	).to(device=lowerCamelCase_	)
				SCREAMING_SNAKE_CASE__  				=					DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name	)
				SCREAMING_SNAKE_CASE__  				=					Features(
				    {"text": Value("string"	), "title": Value("string"	), "embeddings": Sequence(Value("float32"	)	)}	)  # optional, save as float32 instead of float64 to save space
				SCREAMING_SNAKE_CASE__  				=					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
				SCREAMING_SNAKE_CASE__  				=					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
				SCREAMING_SNAKE_CASE__  				=					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
				SCREAMING_SNAKE_CASE__  				=					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
				SCREAMING_SNAKE_CASE__  				=					RagRetriever.from_pretrained(
				    rag_example_args.rag_model_name    ,       index_name="custom"    ,       indexed_dataset=lowerCamelCase_	)
				SCREAMING_SNAKE_CASE__  				=					RagSequenceForGeneration.from_pretrained(rag_example_args.rag_model_name    ,       retriever=lowerCamelCase_	)
				SCREAMING_SNAKE_CASE__  				=					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"	)
				######################################

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


@dataclass
class       lowerCamelCase_		:





				"""simple docstring"""







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


@dataclass
class       lowerCamelCase_		:





				"""simple docstring"""







				_lowerCAmelCase   :   List[Any]   				=       field(
				    default=lowerCamelCase__							,	metadata={
				        "help": "The number of processes to use to split the documents into passages. Default is single process."
				    }							,	)
				_lowerCAmelCase   :   List[Any]   				=       field(
				    default=16							,	metadata={
				        "help": "The batch size to use when computing the passages embeddings using the DPR context encoder."
				    }							,	)


@dataclass
class       lowerCamelCase_		:





				"""simple docstring"""







				_lowerCAmelCase   :   Optional[Any]   				=       field(
				    default=768							,	metadata={"help": "The dimension of the embeddings to pass to the HNSW Faiss index."}							,	)
				_lowerCAmelCase   :   Union[str, Any]   				=       field(
				    default=128							,	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  =	HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
	_lowerCamelCase  =	parser.parse_args_into_dataclasses()
	with TemporaryDirectory() as tmp_dir:
		_lowerCamelCase  =	rag_example_args.output_dir or tmp_dir
		main(rag_example_args, processing_args, index_hnsw_args)



import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         FunnelBaseModel,
         FunnelForMaskedLM,
         FunnelForMultipleChoice,
         FunnelForPreTraining,
         FunnelForQuestionAnswering,
         FunnelForSequenceClassification,
         FunnelForTokenClassification,
         FunnelModel,
     )





class    __UpperCAmelCase	:



   def __init__(				self				:	List[Any],						__A				:	str,						__A				:	Union[str, Any]=1_3,						__A				:	Union[str, Any]=7,						__A				:	List[str]=True,						__A				:	Optional[Any]=True,						__A				:	Tuple=True,						__A				:	Any=True,						__A				:	Dict=9_9,						__A				:	Union[str, Any]=[1, 1, 2],						__A				:	Tuple=1,						__A				:	int=3_2,						__A				:	List[str]=4,						__A				:	List[str]=8,						__A				:	int=3_7,						__A				:	Optional[int]="gelu_new",						__A				:	List[Any]=0.1,						__A				:	str=0.1,						__A				:	int=0.0,						__A				:	Union[str, Any]=5_1_2,						__A				:	int=3,						__A				:	str=0.0_2,						__A				:	Tuple=3,						__A				:	Any=4,						__A				:	Any=None,						__A				:	Any=False,						):
     UpperCAmelCase			:				str										=					parent
     UpperCAmelCase			:				Optional[Any]										=					batch_size
     UpperCAmelCase			:				str										=					seq_length
     UpperCAmelCase			:				str										=					is_training
     UpperCAmelCase			:				Tuple										=					use_input_mask
     UpperCAmelCase			:				Union[str, Any]										=					use_token_type_ids
     UpperCAmelCase			:				Dict										=					use_labels
     UpperCAmelCase			:				Union[str, Any]										=					vocab_size
     UpperCAmelCase			:				Tuple										=					block_sizes
     UpperCAmelCase			:				List[Any]										=					num_decoder_layers
     UpperCAmelCase			:				Union[str, Any]										=					d_model
     UpperCAmelCase			:				Tuple										=					n_head
     UpperCAmelCase			:				Dict										=					d_head
     UpperCAmelCase			:				Optional[int]										=					d_inner
     UpperCAmelCase			:				Optional[int]										=					hidden_act
     UpperCAmelCase			:				int										=					hidden_dropout
     UpperCAmelCase			:				Any										=					attention_dropout
     UpperCAmelCase			:				str										=					activation_dropout
     UpperCAmelCase			:				Dict										=					max_position_embeddings
     UpperCAmelCase			:				Optional[Any]										=					type_vocab_size
     UpperCAmelCase			:				Any										=					2
     UpperCAmelCase			:				List[Any]										=					num_labels
     UpperCAmelCase			:				str										=					num_choices
     UpperCAmelCase			:				str										=					scope
     UpperCAmelCase			:				int										=					initializer_std

     # Used in the tests to check the size of the first attention layer
     UpperCAmelCase			:				int										=					n_head
     # Used in the tests to check the size of the first hidden state
     UpperCAmelCase			:				Dict										=					self.d_model
     # Used in the tests to check the number of output hidden states/attentions
     UpperCAmelCase			:				List[str]										=					sum(self.block_sizes		) + (0 if base else self.num_decoder_layers)
     # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
     # the last hidden state of the first block (which is the first hidden state of the decoder).
     if not base:
       UpperCAmelCase			:				List[str]										=					self.num_hidden_layers + 2



   def 			__magic_name__       (				self				:	List[str]		):
     UpperCAmelCase			:				Dict										=					ids_tensor([self.batch_size, self.seq_length],						self.vocab_size		)

     UpperCAmelCase			:				Optional[int]										=					None
     if self.use_input_mask:
       UpperCAmelCase			:				List[str]										=					ids_tensor([self.batch_size, self.seq_length],						vocab_size=2		)

     UpperCAmelCase			:				int										=					None
     if self.use_token_type_ids:
       UpperCAmelCase			:				Dict										=					ids_tensor([self.batch_size, self.seq_length],						self.type_vocab_size		)

     UpperCAmelCase			:				List[Any]										=					None
     UpperCAmelCase			:				Optional[int]										=					None
     UpperCAmelCase			:				str										=					None
     if self.use_labels:
       UpperCAmelCase			:				List[Any]										=					ids_tensor([self.batch_size],						self.type_sequence_label_size		)
       UpperCAmelCase			:				Optional[int]										=					ids_tensor([self.batch_size, self.seq_length],						self.num_labels		)
       UpperCAmelCase			:				List[str]										=					ids_tensor([self.batch_size],						self.num_choices		)
       UpperCAmelCase			:				Tuple										=					ids_tensor([self.batch_size, self.seq_length],						1		)

     UpperCAmelCase			:				Dict										=					self.get_config()

     return (
         config,
         input_ids,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
         choice_labels,
         fake_token_labels,
     )



   def 			__magic_name__       (				self				:	List[Any]		):
     return FunnelConfig(
         vocab_size=self.vocab_size,						block_sizes=self.block_sizes,						num_decoder_layers=self.num_decoder_layers,						d_model=self.d_model,						n_head=self.n_head,						d_head=self.d_head,						d_inner=self.d_inner,						hidden_act=self.hidden_act,						hidden_dropout=self.hidden_dropout,						attention_dropout=self.attention_dropout,						activation_dropout=self.activation_dropout,						max_position_embeddings=self.max_position_embeddings,						type_vocab_size=self.type_vocab_size,						initializer_std=self.initializer_std,						)



   def 			__magic_name__       (				self				:	Union[str, Any],						__A				:	List[str],						__A				:	Dict,						__A				:	Optional[int],						__A				:	List[Any],						__A				:	List[Any],						__A				:	Optional[int],						__A				:	Dict,						__A				:	Dict,						):
     UpperCAmelCase			:				List[str]										=					FunnelModel(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				List[Any]										=					model(__A,						attention_mask=__A,						token_type_ids=__A		)
     UpperCAmelCase			:				Optional[int]										=					model(__A,						token_type_ids=__A		)
     UpperCAmelCase			:				Optional[int]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, self.seq_length, self.d_model)		)

     UpperCAmelCase			:				Optional[int]										=					False
     UpperCAmelCase			:				List[str]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, self.seq_length, self.d_model)		)

     UpperCAmelCase			:				List[str]										=					False
     UpperCAmelCase			:				Dict										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, self.seq_length, self.d_model)		)



   def 			__magic_name__       (				self				:	Dict,						__A				:	str,						__A				:	List[Any],						__A				:	Union[str, Any],						__A				:	Union[str, Any],						__A				:	Tuple,						__A				:	Tuple,						__A				:	List[Any],						__A				:	Dict,						):
     UpperCAmelCase			:				List[Any]										=					FunnelBaseModel(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				int										=					model(__A,						attention_mask=__A,						token_type_ids=__A		)
     UpperCAmelCase			:				List[Any]										=					model(__A,						token_type_ids=__A		)
     UpperCAmelCase			:				Optional[Any]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, 2, self.d_model)		)

     UpperCAmelCase			:				str										=					False
     UpperCAmelCase			:				Optional[Any]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, 3, self.d_model)		)

     UpperCAmelCase			:				Tuple										=					False
     UpperCAmelCase			:				Optional[Any]										=					model(__A		)
     self.parent.assertEqual(result.last_hidden_state.shape,						(self.batch_size, 2, self.d_model)		)



   def 			__magic_name__       (				self				:	Optional[int],						__A				:	Tuple,						__A				:	Dict,						__A				:	str,						__A				:	Dict,						__A				:	Union[str, Any],						__A				:	List[Any],						__A				:	Any,						__A				:	str,						):
     UpperCAmelCase			:				int										=					self.num_labels
     UpperCAmelCase			:				int										=					FunnelForPreTraining(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				List[Any]										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.seq_length)		)



   def 			__magic_name__       (				self				:	List[str],						__A				:	Optional[int],						__A				:	Union[str, Any],						__A				:	int,						__A				:	Union[str, Any],						__A				:	Optional[int],						__A				:	str,						__A				:	int,						__A				:	str,						):
     UpperCAmelCase			:				List[Any]										=					FunnelForMaskedLM(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Optional[int]										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.seq_length, self.vocab_size)		)



   def 			__magic_name__       (				self				:	Optional[int],						__A				:	Optional[Any],						__A				:	Optional[Any],						__A				:	Dict,						__A				:	Any,						__A				:	Optional[Any],						__A				:	Tuple,						__A				:	Union[str, Any],						__A				:	List[str],						):
     UpperCAmelCase			:				str										=					self.num_labels
     UpperCAmelCase			:				Optional[int]										=					FunnelForSequenceClassification(__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Dict										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.num_labels)		)



   def 			__magic_name__       (				self				:	Union[str, Any],						__A				:	Any,						__A				:	List[str],						__A				:	str,						__A				:	List[Any],						__A				:	List[str],						__A				:	Optional[Any],						__A				:	List[Any],						__A				:	List[str],						):
     UpperCAmelCase			:				str										=					self.num_choices
     UpperCAmelCase			:				Optional[int]										=					FunnelForMultipleChoice(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				List[Any]										=					input_ids.unsqueeze(1		).expand(-1,						self.num_choices,						-1		).contiguous()
     UpperCAmelCase			:				List[Any]										=					token_type_ids.unsqueeze(1		).expand(-1,						self.num_choices,						-1		).contiguous()
     UpperCAmelCase			:				List[Any]										=					input_mask.unsqueeze(1		).expand(-1,						self.num_choices,						-1		).contiguous()
     UpperCAmelCase			:				List[str]										=					model(
         __A,						attention_mask=__A,						token_type_ids=__A,						labels=__A,						)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.num_choices)		)



   def 			__magic_name__       (				self				:	Dict,						__A				:	Tuple,						__A				:	int,						__A				:	Any,						__A				:	Optional[Any],						__A				:	Optional[Any],						__A				:	Tuple,						__A				:	str,						__A				:	List[str],						):
     UpperCAmelCase			:				Union[str, Any]										=					self.num_labels
     UpperCAmelCase			:				Optional[int]										=					FunnelForTokenClassification(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Any										=					model(__A,						attention_mask=__A,						token_type_ids=__A,						labels=__A		)
     self.parent.assertEqual(result.logits.shape,						(self.batch_size, self.seq_length, self.num_labels)		)



   def 			__magic_name__       (				self				:	int,						__A				:	Optional[int],						__A				:	List[str],						__A				:	List[Any],						__A				:	int,						__A				:	str,						__A				:	Any,						__A				:	Union[str, Any],						__A				:	str,						):
     UpperCAmelCase			:				str										=					FunnelForQuestionAnswering(config=__A		)
     model.to(__A		)
     model.eval()
     UpperCAmelCase			:				Optional[Any]										=					model(
         __A,						attention_mask=__A,						token_type_ids=__A,						start_positions=__A,						end_positions=__A,						)
     self.parent.assertEqual(result.start_logits.shape,						(self.batch_size, self.seq_length)		)
     self.parent.assertEqual(result.end_logits.shape,						(self.batch_size, self.seq_length)		)





   def 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase			:				Dict										=					self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		(
         UpperCAmelCase
     )					,		
     )			:				Optional[int]										=					config_and_inputs
     UpperCAmelCase			:				int										=					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
     return config, inputs_dict






@require_torch
class    __UpperCAmelCase	(			lowerCamelCase__     ,    lowerCamelCase__     ,    unittest.TestCase		):
   UpperCamelCase  	=	False
   UpperCamelCase  	=	False
   UpperCamelCase  	=	(
       (
           FunnelModel,
           FunnelForMaskedLM,
           FunnelForPreTraining,
           FunnelForQuestionAnswering,
           FunnelForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   UpperCamelCase  	=	(
       {
           """feature-extraction""": (FunnelBaseModel, FunnelModel),
           """fill-mask""": FunnelForMaskedLM,
           """question-answering""": FunnelForQuestionAnswering,
           """text-classification""": FunnelForSequenceClassification,
           """token-classification""": FunnelForTokenClassification,
           """zero-shot""": FunnelForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )



   def 			__magic_name__       (				self				:	Dict,						__A				:	List[Any],						__A				:	str,						__A				:	Dict=False		):
     UpperCAmelCase			:				int										=					super()._prepare_for_class(__A,						__A,						return_labels=__A		)

     if return_labels:
       if model_class in get_values(__A		):
         UpperCAmelCase			:				Tuple										=					torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length),						dtype=torch.long,						device=__A		)
     return inputs_dict



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				Optional[int]										=					FunnelModelTester(self		)
     UpperCAmelCase			:				str										=					ConfigTester(self,						config_class=__A		)



   def 			__magic_name__       (				self				:	Dict		):
     self.config_tester.run_common_tests()



   def 			__magic_name__       (				self				:	int		):
     UpperCAmelCase			:				Optional[Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A		)



   def 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase			:				List[Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_pretraining(*__A		)



   def 			__magic_name__       (				self				:	Dict		):
     UpperCAmelCase			:				str										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*__A		)



   def 			__magic_name__       (				self				:	str		):
     UpperCAmelCase			:				Any										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*__A		)



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				int										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*__A		)





   def 			__magic_name__       (				self				:	Any,						__A				:	List[Any]		):
     if hasattr(__A,						'''weight'''		) and module.weight is not None:
       module.weight.data.fill_(3		)
     if hasattr(__A,						'''bias'''		) and module.bias is not None:
       module.bias.data.fill_(3		)

     for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
       if hasattr(__A,						__A		) and getattr(__A,						__A		) is not None:
         UpperCAmelCase			:				Optional[int]										=					getattr(__A,						__A		)
         weight.data.fill_(3		)






@require_torch
class    __UpperCAmelCase	(			lowerCamelCase__     ,    unittest.TestCase		):
   UpperCamelCase  	=	False
   UpperCamelCase  	=	False
   UpperCamelCase  	=	(
       (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
   )



   def 			__magic_name__       (				self				:	List[Any]		):
     UpperCAmelCase			:				Optional[int]										=					FunnelModelTester(self,						base=__A		)
     UpperCAmelCase			:				Tuple										=					ConfigTester(self,						config_class=__A		)



   def 			__magic_name__       (				self				:	Union[str, Any]		):
     self.config_tester.run_common_tests()



   def 			__magic_name__       (				self				:	Any		):
     UpperCAmelCase			:				Union[str, Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_base_model(*__A		)



   def 			__magic_name__       (				self				:	Optional[Any]		):
     UpperCAmelCase			:				Union[str, Any]										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*__A		)



   def 			__magic_name__       (				self				:	Optional[int]		):
     UpperCAmelCase			:				Tuple										=					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*__A		)



   def 			__magic_name__       (				self				:	Optional[int]		):
     UpperCAmelCase					,		UpperCAmelCase			:				List[str]										=					self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase			:				List[str]										=					True

     for model_class in self.all_model_classes:
       if model_class.__name__ == "FunnelBaseModel":
         continue
       UpperCAmelCase			:				Any										=					model_class(__A		)
       model.to(__A		)
       model.train()
       UpperCAmelCase			:				str										=					self._prepare_for_class(__A,						__A,						return_labels=__A		)
       UpperCAmelCase			:				str										=					model(**__A		).loss
       loss.backward()





   def 			__magic_name__       (				self				:	str,						__A				:	Any		):
     if hasattr(__A,						'''weight'''		) and module.weight is not None:
       module.weight.data.fill_(3		)
     if hasattr(__A,						'''bias'''		) and module.bias is not None:
       module.bias.data.fill_(3		)

     for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
       if hasattr(__A,						__A		) and getattr(__A,						__A		) is not None:
         UpperCAmelCase			:				Optional[int]										=					getattr(__A,						__A		)
         weight.data.fill_(3		)






@require_torch
@require_sentencepiece
@require_tokenizers
class    __UpperCAmelCase	(			unittest.TestCase		):



   def 			__magic_name__       (				self				:	int		):
     UpperCAmelCase			:				Optional[Any]										=					1_3
     UpperCAmelCase			:				Tuple										=					7
     UpperCAmelCase			:				Optional[Any]										=					torch.arange(0,						batch_size * sequence_length		).long().reshape(__A,						__A		)
     UpperCAmelCase			:				List[Any]										=					[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
     UpperCAmelCase			:				Optional[Any]										=					torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]		)

     UpperCAmelCase			:				Dict										=					FunnelModel.from_pretrained('''sgugger/funnel-random-tiny'''		)
     UpperCAmelCase			:				List[Any]										=					model(__A,						token_type_ids=__A		)[0].abs()

     UpperCAmelCase			:				Union[str, Any]										=					torch.tensor(2_3_4_4.8_3_5_2		)
     UpperCAmelCase			:				List[Any]										=					torch.tensor(0.8_0_5_2		)
     self.assertTrue(torch.allclose(output.sum(),						__A,						atol=1E-4		)		)
     self.assertTrue(torch.allclose(output.mean(),						__A,						atol=1E-4		)		)

     UpperCAmelCase			:				Union[str, Any]										=					torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]		)
     UpperCAmelCase			:				Union[str, Any]										=					model(__A,						attention_mask=__A,						token_type_ids=__A		)[0].abs()

     UpperCAmelCase			:				str										=					torch.tensor(2_3_4_3.8_4_2_5		)
     UpperCAmelCase			:				Optional[Any]										=					torch.tensor(0.8_0_4_9		)
     self.assertTrue(torch.allclose(output.sum(),						__A,						atol=1E-4		)		)
     self.assertTrue(torch.allclose(output.mean(),						__A,						atol=1E-4		)		)





   @slow
   def 			__magic_name__       (				self				:	int		):
     UpperCAmelCase			:				List[Any]										=					FunnelTokenizer.from_pretrained('''huggingface/funnel-small'''		)
     UpperCAmelCase			:				List[str]										=					FunnelModel.from_pretrained('''huggingface/funnel-small'''		)
     UpperCAmelCase			:				List[str]										=					tokenizer('''Hello! I am the Funnel Transformer model.''',						return_tensors='''pt'''		)
     UpperCAmelCase			:				Optional[Any]										=					model(**__A		)[0]

     UpperCAmelCase			:				int										=					torch.tensor(2_3_5.7_2_4_6		)
     UpperCAmelCase			:				Optional[int]										=					torch.tensor(0.0_2_5_6		)
     self.assertTrue(torch.allclose(output.sum(),						__A,						atol=1E-4		)		)
     self.assertTrue(torch.allclose(output.mean(),						__A,						atol=1E-4		)		)





