



import unittest

from transformers import MarkupLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property

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 (
		    MarkupLMForQuestionAnswering,
		    MarkupLMForSequenceClassification,
		    MarkupLMForTokenClassification,
		    MarkupLMModel,
		)

# TODO check dependencies
from transformers import MarkupLMFeatureExtractor, MarkupLMProcessor, MarkupLMTokenizer


class 						lowerCamelCase__			:


						def __init__(			self    :	str   ,							_a    :	List[Any]   ,							_a    :	Optional[int]=1_3   ,							_a    :	List[Any]=7   ,							_a    :	List[str]=True   ,							_a    :	str=True   ,							_a    :	Any=True   ,							_a    :	Optional[Any]=True   ,							_a    :	List[Any]=9_9   ,							_a    :	Any=3_2   ,							_a    :	int=5   ,							_a    :	Dict=4   ,							_a    :	Optional[Any]=3_7   ,							_a    :	str="gelu"   ,							_a    :	Any=0.1   ,							_a    :	List[Any]=0.1   ,							_a    :	str=5_1_2   ,							_a    :	Tuple=1_6   ,							_a    :	Any=2   ,							_a    :	Dict=0.0_2   ,							_a    :	str=3   ,							_a    :	Dict=None   ,							_a    :	Dict=2_0   ,							_a    :	Optional[int]=3_0   ,							_a    :	Any=2   ,							_a    :	List[str]=2   ,							_a    :	int=1_0   ,							):
													a__:							str									=parent
													a__:							Tuple									=batch_size
													a__:							int									=seq_length
													a__:							List[Any]									=is_training
													a__:							Dict									=use_input_mask
													a__:							Union[str, Any]									=use_token_type_ids
													a__:							int									=use_labels
													a__:							List[Any]									=vocab_size
													a__:							List[Any]									=hidden_size
													a__:							Any									=num_hidden_layers
													a__:							Union[str, Any]									=num_attention_heads
													a__:							Optional[int]									=intermediate_size
													a__:							str									=hidden_act
													a__:							str									=hidden_dropout_prob
													a__:							Dict									=attention_probs_dropout_prob
													a__:							Optional[Any]									=max_position_embeddings
													a__:							List[str]									=type_vocab_size
													a__:							int									=type_sequence_label_size
													a__:							Tuple									=initializer_range
													a__:							str									=num_labels
													a__:							Optional[int]									=scope
													a__:							Optional[int]									=max_xpath_tag_unit_embeddings
													a__:							Optional[Any]									=max_xpath_subs_unit_embeddings
													a__:							int									=tag_pad_id
													a__:							Optional[int]									=subs_pad_id
													a__:							Optional[Any]									=max_depth


						def 	_lowerCamelCase      (			self    :	Optional[int]      ):
													a__:							Tuple									=ids_tensor([self.batch_size, self.seq_length]   ,							self.vocab_size      )

													a__:							List[Any]									=ids_tensor(
													    [self.batch_size, self.seq_length, self.max_depth]   ,							self.max_xpath_tag_unit_embeddings      )

													a__:							Optional[Any]									=ids_tensor(
													    [self.batch_size, self.seq_length, self.max_depth]   ,							self.max_xpath_subs_unit_embeddings      )

													a__:							Union[str, Any]									=None
													if self.use_input_mask:
																				a__:							str									=ids_tensor([self.batch_size, self.seq_length]   ,							vocab_size=2      )

													a__:							Optional[Any]									=None
													if self.use_token_type_ids:
																				a__:							Any									=ids_tensor([self.batch_size, self.seq_length]   ,							self.type_vocab_size      )

													a__:							Dict									=None
													a__:							int									=None
													if self.use_labels:
																				a__:							int									=ids_tensor([self.batch_size]   ,							self.type_sequence_label_size      )
																				a__:							List[Any]									=ids_tensor([self.batch_size, self.seq_length]   ,							self.num_labels      )

													a__:							str									=self.get_config()

													return (
													    config,
													    input_ids,
													    xpath_tags_seq,
													    xpath_subs_seq,
													    token_type_ids,
													    input_mask,
													    sequence_labels,
													    token_labels,
													)


						def 	_lowerCamelCase      (			self    :	int      ):
													return MarkupLMConfig(
													    vocab_size=self.vocab_size   ,							hidden_size=self.hidden_size   ,							num_hidden_layers=self.num_hidden_layers   ,							num_attention_heads=self.num_attention_heads   ,							intermediate_size=self.intermediate_size   ,							hidden_act=self.hidden_act   ,							hidden_dropout_prob=self.hidden_dropout_prob   ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,							max_position_embeddings=self.max_position_embeddings   ,							type_vocab_size=self.type_vocab_size   ,							initializer_range=self.initializer_range   ,							max_xpath_tag_unit_embeddings=self.max_xpath_tag_unit_embeddings   ,							max_xpath_subs_unit_embeddings=self.max_xpath_subs_unit_embeddings   ,							tag_pad_id=self.tag_pad_id   ,							subs_pad_id=self.subs_pad_id   ,							max_depth=self.max_depth   ,							)


						def 	_lowerCamelCase      (			self    :	Any   ,							_a    :	Optional[Any]   ,							_a    :	Optional[int]   ,							_a    :	List[str]   ,							_a    :	Dict   ,							_a    :	Optional[int]   ,							_a    :	List[str]   ,							_a    :	Union[str, Any]   ,							_a    :	List[str]   ,							):
													a__:							Any									=MarkupLMModel(config=_a      )
													model.to(_a      )
													model.eval()
													print("Configs:"   ,							model.config.tag_pad_id   ,							model.config.subs_pad_id      )
													a__:							Optional[int]									=model(_a   ,							attention_mask=_a   ,							token_type_ids=_a      )
													a__:							Tuple									=model(_a   ,							token_type_ids=_a      )
													a__:							str									=model(_a      )
													self.parent.assertEqual(result.last_hidden_state.shape   ,							(self.batch_size, self.seq_length, self.hidden_size)      )
													self.parent.assertEqual(result.pooler_output.shape   ,							(self.batch_size, self.hidden_size)      )


						def 	_lowerCamelCase      (			self    :	Dict   ,							_a    :	Optional[int]   ,							_a    :	Optional[int]   ,							_a    :	Any   ,							_a    :	Dict   ,							_a    :	int   ,							_a    :	Optional[int]   ,							_a    :	Optional[Any]   ,							_a    :	Dict   ,							):
													a__:							List[str]									=self.num_labels
													a__:							Union[str, Any]									=MarkupLMForSequenceClassification(_a      )
													model.to(_a      )
													model.eval()
													a__:							Tuple									=model(
													    _a   ,							xpath_tags_seq=_a   ,							xpath_subs_seq=_a   ,							attention_mask=_a   ,							token_type_ids=_a   ,							labels=_a   ,							)
													self.parent.assertEqual(result.logits.shape   ,							(self.batch_size, self.num_labels)      )


						def 	_lowerCamelCase      (			self    :	int   ,							_a    :	Tuple   ,							_a    :	Optional[Any]   ,							_a    :	Optional[int]   ,							_a    :	int   ,							_a    :	Dict   ,							_a    :	str   ,							_a    :	Union[str, Any]   ,							_a    :	Union[str, Any]   ,							):
													a__:							Optional[int]									=self.num_labels
													a__:							int									=MarkupLMForTokenClassification(config=_a      )
													model.to(_a      )
													model.eval()
													a__:							List[str]									=model(
													    _a   ,							xpath_tags_seq=_a   ,							xpath_subs_seq=_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 	_lowerCamelCase      (			self    :	Tuple   ,							_a    :	Tuple   ,							_a    :	int   ,							_a    :	List[str]   ,							_a    :	Any   ,							_a    :	str   ,							_a    :	Union[str, Any]   ,							_a    :	Union[str, Any]   ,							_a    :	List[str]   ,							):
													a__:							List[str]									=MarkupLMForQuestionAnswering(config=_a      )
													model.to(_a      )
													model.eval()
													a__:							int									=model(
													    _a   ,							xpath_tags_seq=_a   ,							xpath_subs_seq=_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 	_lowerCamelCase      (			self    :	Optional[int]      ):
													a__:							List[str]									=self.prepare_config_and_inputs()
													(
													    (
													    a__
													)				,      (
													    a__
													)				,      (
													    a__
													)				,      (
													    a__
													)				,      (
													    a__
													)				,      (
													    a__
													)				,      (
													    a__
													)				,      (
													    a__
													)				,      
													):							Optional[int]									=config_and_inputs
													a__:							Optional[int]									={
													    "input_ids": input_ids,
													    "xpath_tags_seq": xpath_tags_seq,
													    "xpath_subs_seq": xpath_subs_seq,
													    "token_type_ids": token_type_ids,
													    "attention_mask": input_mask,
													}
													return config, inputs_dict


@require_torch
class 						lowerCamelCase__			(		_a					,       _a					,       unittest.TestCase ):
						_lowerCAmelCase             =			(
						    (
						        MarkupLMModel,
						        MarkupLMForSequenceClassification,
						        MarkupLMForTokenClassification,
						        MarkupLMForQuestionAnswering,
						    )
						    if is_torch_available()
						    else None
						)
						_lowerCAmelCase             =			(
						    {
						        '''feature-extraction''': MarkupLMModel,
						        '''question-answering''': MarkupLMForQuestionAnswering,
						        '''text-classification''': MarkupLMForSequenceClassification,
						        '''token-classification''': MarkupLMForTokenClassification,
						        '''zero-shot''': MarkupLMForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)


						def 	_lowerCamelCase      (			self    :	Dict   ,							_a    :	int   ,							_a    :	List[Any]   ,							_a    :	Optional[Any]   ,							_a    :	Optional[int]   ,							_a    :	List[Any]      ):
													# ValueError: Nodes must be of type `List[str]` (single pretokenized example), or `List[List[str]]`
													# (batch of pretokenized examples).
													return True


						def 	_lowerCamelCase      (			self    :	int      ):
													a__:							List[Any]									=MarkupLMModelTester(self      )
													a__:							Dict									=ConfigTester(self   ,							config_class=_a   ,							hidden_size=3_7      )


						def 	_lowerCamelCase      (			self    :	int      ):
													self.config_tester.run_common_tests()


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							Optional[Any]									=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*_a      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							List[str]									=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_sequence_classification(*_a      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							List[Any]									=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*_a      )







						def 	_lowerCamelCase      (			self    :	Optional[Any]      ):
													a__:							Any									=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_question_answering(*_a      )





def     __lowerCamelCase							(   ):
							a__:							Union[str, Any]									="\n    <!DOCTYPE html>\n    <html>\n    <head>\n    <title>Page Title</title>\n    </head>\n    <body>\n\n    <h1>This is a Heading</h1>\n    <p>This is a paragraph.</p>\n\n    </body>\n    </html>\n    "

							return html_string







@require_torch
class 						lowerCamelCase__			(		unittest.TestCase ):


						@cached_property
						def 	_lowerCamelCase      (			self    :	int      ):
													# TODO use from_pretrained here
													a__:							Dict									=MarkupLMFeatureExtractor()
													a__:							int									=MarkupLMTokenizer.from_pretrained("microsoft/markuplm-base"      )

													return MarkupLMProcessor(_a   ,							_a      )







						@slow
						def 	_lowerCamelCase      (			self    :	Optional[int]      ):
													a__:							Tuple									=MarkupLMModel.from_pretrained("microsoft/markuplm-base"      ).to(_a      )

													a__:							Union[str, Any]									=self.default_processor

													a__:							Optional[int]									=processor(prepare_html_string()   ,							return_tensors="pt"      )
													a__:							Any									=inputs.to(_a      )

													# forward pass
													with torch.no_grad():
																				a__:							Tuple									=model(**_a      )

													# verify the last hidden states
													a__:							str									=torch.Size([1, 1_4, 7_6_8]      )
													self.assertEqual(outputs.last_hidden_state.shape   ,							_a      )

													a__:							Tuple									=torch.tensor(
													    [[0.0_6_7_5, -0.0_0_5_2, 0.5_0_0_1], [-0.2_2_8_1, 0.0_8_0_2, 0.2_1_9_2], [-0.0_5_8_3, -0.3_3_1_1, 0.1_1_8_5]]      ).to(_a      )

													self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]   ,							_a   ,							atol=1e-4      )      )










import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


__UpperCAmelCase    	=  logging.getLogger(__name__)


class 						lowerCamelCase__			:


						def __init__(			self    :	Optional[Any]      ):
													a__:							Tuple									=False


						def 	_lowerCamelCase      (			self    :	Dict   ,							_a    :	Any   ,							_a    :	Union[str, Any]   ,							_a    :	List[str]   ,							_a    :	Optional[int]      ):
													if not self.initialized:
																				a__:							Any									=RagRetriever(
																				    _a   ,							question_encoder_tokenizer=_a   ,							generator_tokenizer=_a   ,							index=_a   ,							init_retrieval=_a   ,							)
																				a__:							str									=True


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													self.retriever.index.init_index()


						def 	_lowerCamelCase      (			self    :	Optional[int]      ):
													# delete the old self.retriever object before assigning the new index
													del self.retriever
													a__:							Dict									=False







						def 	_lowerCamelCase      (			self    :	str   ,							_a    :	List[Any]   ,							_a    :	Any      ):
													a__				,      a__:							List[str]									=self.retriever._main_retrieve(_a   ,							_a      )
													a__:							List[str]									=self.retriever.index.get_doc_dicts(_a      )
													return doc_ids, retrieved_doc_embeds, doc_dicts


class 						lowerCamelCase__			(		_a ):


						def __init__(			self    :	Optional[int]   ,							_a    :	List[Any]   ,							_a    :	Any   ,							_a    :	Optional[Any]   ,							_a    :	Union[str, Any]   ,							_a    :	str=None      ):
													if index is not None and index.is_initialized() and len(_a      ) > 0:
																				raise ValueError(
																				    "When using Ray for distributed fine-tuning, "
																				    "you'll need to provide the paths instead, "
																				    "as the dataset and the index are loaded "
																				    "separately. More info in examples/rag/use_own_knowledge_dataset.py "      )

													super().__init__(
													    _a   ,							question_encoder_tokenizer=_a   ,							generator_tokenizer=_a   ,							index=_a   ,							init_retrieval=_a   ,							)

													a__:							Optional[Any]									=retrieval_workers
													a__:							str									=question_encoder_tokenizer
													a__:							Optional[int]									=generator_tokenizer
													if len(self.retrieval_workers      ) > 0:
																				ray.get(
																				    [
																				        worker.create_rag_retriever.remote(_a   ,							_a   ,							_a   ,							_a      )
																				        for worker in self.retrieval_workers
																				    ]      )


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													logger.info("initializing retrieval"      )

													if len(self.retrieval_workers      ) > 0:
																				ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]      )
													else:
																				# Non-distributed training. Load index into this same process.
																				self.index.init_index()


						def 	_lowerCamelCase      (			self    :	List[Any]   ,							_a    :	Union[str, Any]   ,							_a    :	Union[str, Any]      ):
													if len(self.retrieval_workers      ) > 0:
																				# Select a random retrieval actor.
																				a__:							Tuple									=self.retrieval_workers[random.randint(0   ,							len(self.retrieval_workers      ) - 1      )]
																				a__				,      a__				,      a__:							Union[str, Any]									=ray.get(
																				    random_worker.retrieve.remote(_a   ,							_a      )      )
													else:
																				a__				,      a__:							List[Any]									=self._main_retrieve(_a   ,							_a      )
																				a__:							List[str]									=self.index.get_doc_dicts(_a      )
													return retrieved_doc_embeds, doc_ids, doc_dicts


						@classmethod
						def 	_lowerCamelCase      (			cls    :	str   ,							_a    :	Union[str, Any]   ,							_a    :	Any=None   ,							**_a    :	Dict      ):
													return super(_a   ,							cls      ).get_tokenizers(_a   ,							_a   ,							**_a      )


						@classmethod
						def 	_lowerCamelCase      (			cls    :	Dict   ,							_a    :	Optional[int]   ,							_a    :	Tuple   ,							_a    :	int=None   ,							**_a    :	List[str]      ):
													a__:							Union[str, Any]									=kwargs.pop("config"   ,							_a      ) or RagConfig.from_pretrained(_a   ,							**_a      )
													a__:							List[Any]									=RagTokenizer.from_pretrained(_a   ,							config=_a      )
													a__:							str									=rag_tokenizer.question_encoder
													a__:							Any									=rag_tokenizer.generator

													if indexed_dataset is not None:
																				a__:							Dict									="custom"
																				a__:							Dict									=CustomHFIndex(config.retrieval_vector_size   ,							_a      )
													else:
																				a__:							List[str]									=cls._build_index(_a      )

													return cls(
													    _a   ,							question_encoder_tokenizer=_a   ,							generator_tokenizer=_a   ,							retrieval_workers=_a   ,							index=_a   ,							)







						def 	_lowerCamelCase      (			self    :	Union[str, Any]      ):
													logger.info("re-loading the new dataset with embeddings"      )
													# access from the training loop

													ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]      )

													# build the index object again
													a__:							str									=self._build_index(self.config      )

													ray.get(
													    [
													        worker.create_rag_retriever.remote(
													            self.config   ,							self.question_encoder_tokenizer   ,							self.generator_tokenizer   ,							_a      )
													        for worker in self.retrieval_workers
													    ]      )





