



import os
from functools import partial
from glob import glob

import faiss
from datasets import Features, Sequence, Value, concatenate_datasets, load_dataset, load_from_disk

from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast
def 			UpperCamelCase__						(    A__      ,			A__=100      ,			A__=" "   )       -> Any:
				snake_case__ :     Union[str, Any] 							=				text.split(A__   )
				return [character.join(text[i : i + n]   ).strip() for i in range(0      ,			len(A__   )      ,			A__   )]
def 			UpperCamelCase__						(    A__   )       -> Union[str, Any]:
				snake_case__ :     List[str] 							=				[], []
				for title, text in zip(documents['title']      ,			documents['text']   ):
								if text is not None:
												for passage in split_text(A__   ):
																titles.append(title if title is not None else ''   )
																texts.append(A__   )
				return {"title": titles, "text": texts}
def 			UpperCamelCase__						(    A__      ,			A__      ,			A__      ,			A__      ,			A__      ,			A__   )       -> Tuple:
				snake_case__ :     Optional[Any] 							=				load_dataset(
				    'csv'      ,			data_files=[csv_path]      ,			split='train'      ,			delimiter='\t'      ,			column_names=['title', 'text']   )
				snake_case__ :     str 							=				kb_dataset.map(
				    A__      ,			batched=A__      ,			num_proc=1   )  # if you want you can load already splitted csv.
				snake_case__ :     Dict 							=				[kb_dataset.shard(A__      ,			A__      ,			contiguous=A__   ) for i in range(A__   )]
				snake_case__ :     List[str] 							=				kb_list[process_num]

				snake_case__ :     List[str] 							=				'data_' + str(A__   )
				snake_case__ :     Optional[Any] 							=				os.path.join(A__      ,			A__   )

				snake_case__ :     int 							=				DPRContextEncoderTokenizerFast.from_pretrained('facebook/dpr-ctx_encoder-multiset-base'   )
				snake_case__ :     Any 							=				ctx_encoder.to(device=A__   )

				def embed(
				    A__      ,			A__      ,			A__      ,			A__   ) -> dict:
								snake_case__ :     Optional[int] 							=				ctx_tokenizer(
								    documents['title']      ,			documents['text']      ,			truncation=A__      ,			padding='longest'      ,			return_tensors='pt'   )['input_ids']
								snake_case__ :     str 							=				ctx_encoder(input_ids.to(device=A__   )      ,			return_dict=A__   ).pooler_output
								return {"embeddings": embeddings.detach().cpu().numpy()}

				snake_case__ :     Union[str, Any] 							=				Features(
				    {'text': Value('string'   ), 'title': Value('string'   ), 'embeddings': Sequence(Value('float32'   )   )}   )  # optional, save as float32 instead of float64 to save space

				snake_case__ :     Optional[Any] 							=				data_shrad.map(
				    partial(A__      ,			ctx_encoder=A__      ,			ctx_tokenizer=A__      ,			device=A__   )      ,			batched=A__      ,			batch_size=16      ,			features=A__      ,			)
				dataset.save_to_disk(A__   )

def 			UpperCamelCase__						(    A__      ,			A__   )       -> Optional[int]:
				snake_case__ :     int 							=				[]

				for shard_address in glob(str(A__   ) + '/*/'   ):
								data_shard_list.append(load_from_disk(A__   )   )

				snake_case__ :     Union[str, Any] 							=				concatenate_datasets(A__   )
				faiss.omp_set_num_threads(96   )

				snake_case__ :     str 							=				faiss.IndexHNSWFlat(768      ,			128      ,			faiss.METRIC_INNER_PRODUCT   )
				concat.add_faiss_index('embeddings'      ,			custom_index=A__   )
				concat.get_index('embeddings'   ).save(
				    A__   )  # since we load the index in to memory,we can directly update the index in the disk



from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.openai.modeling_tf_openai import (
				    TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFOpenAIGPTDoubleHeadsModel,
				    TFOpenAIGPTForSequenceClassification,
				    TFOpenAIGPTLMHeadModel,
				    TFOpenAIGPTModel,
				)





class 			__snake_case  :




					def __init__(       self   ,	__UpperCamelCase   ,	)					->  Tuple:



									'''simple docstring'''







									snake_case__ :     List[str] 							=				parent
									snake_case__ :     Optional[int] 							=				13
									snake_case__ :     Tuple 							=				7
									snake_case__ :     Optional[int] 							=				True
									snake_case__ :     int 							=				True
									snake_case__ :     int 							=				True
									snake_case__ :     int 							=				True
									snake_case__ :     Union[str, Any] 							=				True
									snake_case__ :     Optional[Any] 							=				99
									snake_case__ :     Dict 							=				32
									snake_case__ :     Union[str, Any] 							=				2
									snake_case__ :     Tuple 							=				4
									snake_case__ :     List[str] 							=				37
									snake_case__ :     Optional[int] 							=				'gelu'
									snake_case__ :     List[str] 							=				0.1
									snake_case__ :     Union[str, Any] 							=				0.1
									snake_case__ :     str 							=				512
									snake_case__ :     str 							=				16
									snake_case__ :     Dict 							=				2
									snake_case__ :     List[Any] 							=				0.0_2
									snake_case__ :     Union[str, Any] 							=				3
									snake_case__ :     Optional[Any] 							=				4
									snake_case__ :     List[Any] 							=				None
									snake_case__ :     Union[str, Any] 							=				self.vocab_size - 1




					def       __a      (       self		)					->  Dict:



									'''simple docstring'''







									snake_case__ :     Any 							=				ids_tensor([self.batch_size, self.seq_length]   ,	self.vocab_size		)

									snake_case__ :     List[Any] 							=				None
									if self.use_input_mask:
													snake_case__ :     str 							=				random_attention_mask([self.batch_size, self.seq_length]		)

									snake_case__ :     List[str] 							=				None
									if self.use_token_type_ids:
													snake_case__ :     List[Any] 							=				ids_tensor([self.batch_size, self.seq_length]   ,	self.type_vocab_size		)

									snake_case__ :     Dict 							=				None
									if self.use_mc_token_ids:
													snake_case__ :     Any 							=				ids_tensor([self.batch_size, self.num_choices]   ,	self.seq_length		)

									snake_case__ :     List[Any] 							=				None
									snake_case__ :     Dict 							=				None
									snake_case__ :     str 							=				None
									if self.use_labels:
													snake_case__ :     Tuple 							=				ids_tensor([self.batch_size]   ,	self.type_sequence_label_size		)
													snake_case__ :     Tuple 							=				ids_tensor([self.batch_size, self.seq_length]   ,	self.num_labels		)
													snake_case__ :     Tuple 							=				ids_tensor([self.batch_size]   ,	self.num_choices		)

									snake_case__ :     str 							=				OpenAIGPTConfig(
									    vocab_size=self.vocab_size   ,	n_embd=self.hidden_size   ,	n_layer=self.num_hidden_layers   ,	n_head=self.num_attention_heads   ,	n_positions=self.max_position_embeddings   ,	pad_token_id=self.pad_token_id   ,	)

									snake_case__ :     Union[str, Any] 							=				ids_tensor([self.num_hidden_layers, self.num_attention_heads]   ,	2		)

									return (
									    config,
									    input_ids,
									    input_mask,
									    head_mask,
									    token_type_ids,
									    mc_token_ids,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									)




					def       __a      (       self   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	*__UpperCamelCase		)					->  List[str]:



									'''simple docstring'''







									snake_case__ :     List[Any] 							=				TFOpenAIGPTModel(config=__UpperCamelCase		)
									snake_case__ :     str 							=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									snake_case__ :     Dict 							=				model(__UpperCamelCase		)

									snake_case__ :     Tuple 							=				[input_ids, input_mask]
									snake_case__ :     List[str] 							=				model(__UpperCamelCase		)

									snake_case__ :     Tuple 							=				model(__UpperCamelCase		)

									self.parent.assertEqual(result.last_hidden_state.shape   ,	(self.batch_size, self.seq_length, self.hidden_size)		)




					def       __a      (       self   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	*__UpperCamelCase		)					->  Optional[int]:



									'''simple docstring'''







									snake_case__ :     List[str] 							=				TFOpenAIGPTLMHeadModel(config=__UpperCamelCase		)
									snake_case__ :     Optional[int] 							=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									snake_case__ :     List[str] 							=				model(__UpperCamelCase		)
									self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.seq_length, self.vocab_size)		)




					def       __a      (       self   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	*__UpperCamelCase		)					->  Tuple:



									'''simple docstring'''







									snake_case__ :     Dict 							=				TFOpenAIGPTDoubleHeadsModel(config=__UpperCamelCase		)

									snake_case__ :     Optional[int] 							=				tf.tile(tf.expand_dims(__UpperCamelCase   ,	1		)   ,	(1, self.num_choices, 1)		)
									snake_case__ :     Tuple 							=				tf.tile(tf.expand_dims(__UpperCamelCase   ,	1		)   ,	(1, self.num_choices, 1)		)
									snake_case__ :     Any 							=				tf.tile(tf.expand_dims(__UpperCamelCase   ,	1		)   ,	(1, self.num_choices, 1)		)

									snake_case__ :     Dict 							=				{
									    'input_ids': multiple_choice_inputs_ids,
									    'mc_token_ids': mc_token_ids,
									    'attention_mask': multiple_choice_input_mask,
									    'token_type_ids': multiple_choice_token_type_ids,
									}
									snake_case__ :     Tuple 							=				model(__UpperCamelCase		)
									self.parent.assertEqual(
									    result.logits.shape   ,	(self.batch_size, self.num_choices, self.seq_length, self.vocab_size)		)
									self.parent.assertEqual(result.mc_logits.shape   ,	(self.batch_size, self.num_choices)		)




					def       __a      (       self   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	*__UpperCamelCase		)					->  Tuple:



									'''simple docstring'''







									snake_case__ :     int 							=				self.num_labels
									snake_case__ :     int 							=				ids_tensor([self.batch_size]   ,	self.type_sequence_label_size		)
									snake_case__ :     str 							=				{
									    'input_ids': input_ids,
									    'attention_mask': input_mask,
									    'token_type_ids': token_type_ids,
									    'labels': sequence_labels,
									}
									snake_case__ :     str 							=				TFOpenAIGPTForSequenceClassification(__UpperCamelCase		)
									snake_case__ :     int 							=				model(__UpperCamelCase		)
									self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.num_labels)		)




					def       __a      (       self		)					->  List[str]:



									'''simple docstring'''







									snake_case__ :     Dict 							=				self.prepare_config_and_inputs()

									(
									    (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  (
									    snake_case__
									)							,  
									) :     Any 							=				config_and_inputs

									snake_case__ :     List[str] 							=				{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
									return config, inputs_dict





@require_tf
class 			__snake_case  (      _lowerCamelCase   ,_lowerCamelCase   ,unittest.TestCase     ):
					__lowerCamelCase	 =       (
					    (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
					    if is_tf_available()
					    else ()
					)
					__lowerCamelCase	 =       (
					    (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
					)  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
					__lowerCamelCase	 =       (
					    {
					        """feature-extraction""": TFOpenAIGPTModel,
					        """text-classification""": TFOpenAIGPTForSequenceClassification,
					        """text-generation""": TFOpenAIGPTLMHeadModel,
					        """zero-shot""": TFOpenAIGPTForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					__lowerCamelCase	 =       False
					__lowerCamelCase	 =       False




					def       __a      (       self   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase   ,	__UpperCamelCase		)					->  int:



									'''simple docstring'''







									if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
													# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
													# `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
													# tiny config could not be created.
													return True

									return False




					def       __a      (       self		)					->  Any:



									'''simple docstring'''







									snake_case__ :     Tuple 							=				TFOpenAIGPTModelTester(self		)
									snake_case__ :     Any 							=				ConfigTester(self   ,	config_class=__UpperCamelCase   ,	n_embd=37		)




					def       __a      (       self		)					->  int:



									'''simple docstring'''







									self.config_tester.run_common_tests()




					def       __a      (       self		)					->  Dict:



									'''simple docstring'''







									snake_case__ :     Optional[Any] 							=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_model(*__UpperCamelCase		)




					def       __a      (       self		)					->  int:



									'''simple docstring'''







									snake_case__ :     str 							=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_lm_head(*__UpperCamelCase		)




					def       __a      (       self		)					->  List[Any]:



									'''simple docstring'''







									snake_case__ :     Dict 							=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_double_head(*__UpperCamelCase		)




					def       __a      (       self		)					->  Dict:



									'''simple docstring'''







									snake_case__ :     str 							=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*__UpperCamelCase		)




					@slow
					def       __a      (       self		)					->  Union[str, Any]:



									'''simple docstring'''







									for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													snake_case__ :     Optional[Any] 							=				TFOpenAIGPTModel.from_pretrained(__UpperCamelCase		)
													self.assertIsNotNone(__UpperCamelCase		)





@require_tf
class 			__snake_case  (      unittest.TestCase     ):




					@slow
					def       __a      (       self		)					->  Union[str, Any]:



									'''simple docstring'''







									snake_case__ :     Dict 							=				TFOpenAIGPTLMHeadModel.from_pretrained('openai-gpt'		)
									snake_case__ :     Any 							=				tf.convert_to_tensor([[481, 4735, 544]]   ,	dtype=tf.intaa		)  # the president is
									snake_case__ :     int 							=				[
									    481,
									    4735,
									    544,
									    246,
									    963,
									    870,
									    762,
									    239,
									    244,
									    40477,
									    244,
									    249,
									    719,
									    881,
									    487,
									    544,
									    240,
									    244,
									    603,
									    481,
									]  # the president is a very good man. " \n " i\'m sure he is, " said the

									snake_case__ :     int 							=				model.generate(__UpperCamelCase   ,	do_sample=__UpperCamelCase		)
									self.assertListEqual(output_ids[0].numpy().tolist()   ,	__UpperCamelCase		)


