





import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
					import torch
					from torch import nn

					from transformers import (
					    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
					    TimesformerForVideoClassification,
					    TimesformerModel,
					)
					from transformers.models.timesformer.modeling_timesformer import TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from transformers import VideoMAEImageProcessor


class     A     :
			'''simple docstring'''




			def __init__(     self  :       Tuple				,      __lowerCAmelCase  :       int				,      __lowerCAmelCase  :       List[Any]=13				,      __lowerCAmelCase  :       str=10				,      __lowerCAmelCase  :       List[Any]=3				,      __lowerCAmelCase  :       Optional[Any]=2				,      __lowerCAmelCase  :       Optional[Any]=2				,      __lowerCAmelCase  :       Optional[int]=True				,      __lowerCAmelCase  :       Tuple=True				,      __lowerCAmelCase  :       List[Any]=32				,      __lowerCAmelCase  :       str=5				,      __lowerCAmelCase  :       Optional[Any]=4				,      __lowerCAmelCase  :       Optional[Any]=37				,      __lowerCAmelCase  :       Any="gelu"				,      __lowerCAmelCase  :       List[str]=0.1				,      __lowerCAmelCase  :       List[str]=0.1				,      __lowerCAmelCase  :       List[str]=10				,      __lowerCAmelCase  :       Dict=0.0_2				,      __lowerCAmelCase  :       Tuple="divided_space_time"				,      __lowerCAmelCase  :       Dict=None				,      )      ->      Optional[int]:







					"""simple docstring"""






					A__     =  parent
					A__     =  batch_size
					A__     =  image_size
					A__     =  num_channels
					A__     =  patch_size
					A__     =  num_frames
					A__     =  is_training
					A__     =  use_labels
					A__     =  hidden_size
					A__     =  num_hidden_layers
					A__     =  num_attention_heads
					A__     =  intermediate_size
					A__     =  hidden_act
					A__     =  hidden_dropout_prob
					A__     =  attention_probs_dropout_prob
					A__     =  attention_type
					A__     =  initializer_range
					A__     =  scope
					A__     =  num_labels

					# in TimeSformer, the number of spatial tokens equals num_frames * num_patches per frame + 1 CLS token
					A__     =  (image_size // patch_size) ** 2
					A__     =  (num_frames) * self.num_patches_per_frame + 1






			def   a_			(     self  :       List[Any]  )      ->      List[str]:







					"""simple docstring"""






					A__     =  floats_tensor(
					    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]  )

					A__     =  None
					if self.use_labels:
							A__     =  ids_tensor([self.batch_size]				,      self.num_labels  )

					A__     =  self.get_config()

					return config, pixel_values, labels






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







					"""simple docstring"""






					A__     =  TimesformerConfig(
					    image_size=self.image_size				,      patch_size=self.patch_size				,      num_channels=self.num_channels				,      num_frames=self.num_frames				,      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				,      initializer_range=self.initializer_range				,      attention_type=self.attention_type				,      )
					A__     =  self.num_labels
					return config






			def   a_			(     self  :       str				,      __lowerCAmelCase  :       Union[str, Any]				,      __lowerCAmelCase  :       Any				,      __lowerCAmelCase  :       Union[str, Any]  )      ->      Tuple:







					"""simple docstring"""






					A__     =  TimesformerModel(config=__lowerCAmelCase  )
					model.to(__lowerCAmelCase  )
					model.eval()
					A__     =  model(__lowerCAmelCase  )
					self.parent.assertEqual(result.last_hidden_state.shape				,      (self.batch_size, self.seq_length, self.hidden_size)  )






			def   a_			(     self  :       Optional[int]				,      __lowerCAmelCase  :       str				,      __lowerCAmelCase  :       Any				,      __lowerCAmelCase  :       Optional[int]  )      ->      Tuple:







					"""simple docstring"""






					A__     =  TimesformerForVideoClassification(__lowerCAmelCase  )
					model.to(__lowerCAmelCase  )
					model.eval()

					A__     =  model(__lowerCAmelCase  )

					# verify the logits shape
					A__     =  torch.Size((self.batch_size, self.num_labels)  )
					self.parent.assertEqual(result.logits.shape				,      __lowerCAmelCase  )







			def   a_			(     self  :       int  )      ->      Optional[int]:







					"""simple docstring"""






					A__     =  self.prepare_config_and_inputs()
					A__ ,     A__ ,     A__     =  config_and_inputs
					A__     =  {"""pixel_values""": pixel_values}
					return config, inputs_dict


@require_torch
class     A     (SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE  ,	unittest.TestCase       ):
			'''simple docstring'''




			__lowerCamelCase      :			Dict							     =       (TimesformerModel, TimesformerForVideoClassification) if is_torch_available() else ()
			__lowerCamelCase      :			List[str]							     =       (
			    {'''feature-extraction''': TimesformerModel, '''video-classification''': TimesformerForVideoClassification}
			    if is_torch_available()
			    else {}
			)

			__lowerCamelCase      :			Union[str, Any]							     =       False
			__lowerCamelCase      :			Tuple							     =       False
			__lowerCamelCase      :			str							     =       False
			__lowerCamelCase      :			Any							     =       False






			def   a_			(     self  :       int  )      ->      List[Any]:







					"""simple docstring"""






					A__     =  TimesformerModelTester(self  )
					A__     =  ConfigTester(
					    self				,      config_class=__lowerCAmelCase				,      has_text_modality=__lowerCAmelCase				,      hidden_size=37  )






			def   a_			(     self  :       Any				,      __lowerCAmelCase  :       Dict				,      __lowerCAmelCase  :       List[str]				,      __lowerCAmelCase  :       Dict=False  )      ->      str:







					"""simple docstring"""






					A__     =  copy.deepcopy(__lowerCAmelCase  )

					if return_labels:
							if model_class in get_values(__lowerCAmelCase  ):
									A__     =  torch.zeros(
									    self.model_tester.batch_size				,      dtype=torch.long				,      device=__lowerCAmelCase  )

					return inputs_dict






			def   a_			(     self  :       Union[str, Any]  )      ->      Optional[Any]:







					"""simple docstring"""






					self.config_tester.run_common_tests()






			@unittest.skip(reason="""TimeSformer does not use inputs_embeds"""  )
			def   a_			(     self  :       Union[str, Any]  )      ->      Optional[Any]:







					"""simple docstring"""






					pass






			def   a_			(     self  :       Tuple  )      ->      Any:







					"""simple docstring"""






					A__ ,     A__     =  self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							A__     =  model_class(__lowerCAmelCase  )
							self.assertIsInstance(model.get_input_embeddings()				,      (nn.Module)  )
							A__     =  model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(__lowerCAmelCase				,      nn.Linear  )  )






			def   a_			(     self  :       List[str]  )      ->      int:







					"""simple docstring"""






					A__ ,     A__     =  self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							A__     =  model_class(__lowerCAmelCase  )
							A__     =  inspect.signature(model.forward  )
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							A__     =  [*signature.parameters.keys()]

							A__     =  ["""pixel_values"""]
							self.assertListEqual(arg_names[:1]				,      __lowerCAmelCase  )






			def   a_			(     self  :       List[Any]  )      ->      Union[str, Any]:







					"""simple docstring"""






					A__     =  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*__lowerCAmelCase  )






			def   a_			(     self  :       int  )      ->      Any:







					"""simple docstring"""






					A__     =  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_video_classification(*__lowerCAmelCase  )






			@slow
			def   a_			(     self  :       int  )      ->      Dict:







					"""simple docstring"""






					for model_name in TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__     =  TimesformerModel.from_pretrained(__lowerCAmelCase  )
							self.assertIsNotNone(__lowerCAmelCase  )






			def   a_			(     self  :       Optional[int]  )      ->      Optional[Any]:







					"""simple docstring"""






					if not self.has_attentions:
							pass

					else:
							A__ ,     A__     =  self.model_tester.prepare_config_and_inputs_for_common()
							A__     =  True

							for model_class in self.all_model_classes:
									A__     =  self.model_tester.seq_length
									A__     =  self.model_tester.num_frames

									A__     =  True
									A__     =  False
									A__     =  True
									A__     =  model_class(__lowerCAmelCase  )
									model.to(__lowerCAmelCase  )
									model.eval()
									with torch.no_grad():
											A__     =  model(**self._prepare_for_class(__lowerCAmelCase				,      __lowerCAmelCase  )  )
									A__     =  outputs.attentions
									self.assertEqual(len(__lowerCAmelCase  )				,      self.model_tester.num_hidden_layers  )

									# check that output_attentions also work using config
									del inputs_dict["output_attentions"]
									A__     =  True
									A__     =  model_class(__lowerCAmelCase  )
									model.to(__lowerCAmelCase  )
									model.eval()
									with torch.no_grad():
											A__     =  model(**self._prepare_for_class(__lowerCAmelCase				,      __lowerCAmelCase  )  )
									A__     =  outputs.attentions
									self.assertEqual(len(__lowerCAmelCase  )				,      self.model_tester.num_hidden_layers  )

									# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
									self.assertListEqual(
									    list(attentions[0].shape[-3:]  )				,      [self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]				,      )
									A__     =  len(__lowerCAmelCase  )

									# Check attention is always last and order is fine
									A__     =  True
									A__     =  True
									A__     =  model_class(__lowerCAmelCase  )
									model.to(__lowerCAmelCase  )
									model.eval()
									with torch.no_grad():
											A__     =  model(**self._prepare_for_class(__lowerCAmelCase				,      __lowerCAmelCase  )  )

									self.assertEqual(out_len + 1				,      len(__lowerCAmelCase  )  )

									A__     =  outputs.attentions

									self.assertEqual(len(__lowerCAmelCase  )				,      self.model_tester.num_hidden_layers  )

									# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
									self.assertListEqual(
									    list(self_attentions[0].shape[-3:]  )				,      [self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]				,      )







			def   a_			(     self  :       Dict  )      ->      Tuple:







					"""simple docstring"""






					def check_hidden_states_output(__lowerCAmelCase  :       List[str]				,      __lowerCAmelCase  :       str				,      __lowerCAmelCase  :       Dict  ):
							A__     =  model_class(__lowerCAmelCase  )
							model.to(__lowerCAmelCase  )
							model.eval()

							with torch.no_grad():
									A__     =  model(**self._prepare_for_class(__lowerCAmelCase				,      __lowerCAmelCase  )  )

							A__     =  outputs.hidden_states
							A__     =  self.model_tester.num_hidden_layers + 1
							self.assertEqual(len(__lowerCAmelCase  )				,      __lowerCAmelCase  )

							A__     =  self.model_tester.seq_length

							self.assertListEqual(
							    list(hidden_states[0].shape[-2:]  )				,      [seq_length, self.model_tester.hidden_size]				,      )

					A__ ,     A__     =  self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							A__     =  True
							check_hidden_states_output(__lowerCAmelCase				,      __lowerCAmelCase				,      __lowerCAmelCase  )

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							A__     =  True

							check_hidden_states_output(__lowerCAmelCase				,      __lowerCAmelCase				,      __lowerCAmelCase  )







def   __lowerCamelCase   (				)						->							Any:






		"""simple docstring"""

		A__     =  hf_hub_download(
		    repo_id="""hf-internal-testing/spaghetti-video"""       ,	filename="""eating_spaghetti.npy"""       ,	repo_type="""dataset"""       )
		A__     =  np.load(__a       )
		return list(__a       )






@require_torch
@require_vision
class     A     (unittest.TestCase       ):
			'''simple docstring'''




			@cached_property
			def   a_			(     self  :       Union[str, Any]  )      ->      Any:







					"""simple docstring"""






					return (
					    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]				,      image_std=[0.5, 0.5, 0.5]  )
					    if is_vision_available()
					    else None
					)







			@slow
			def   a_			(     self  :       str  )      ->      Tuple:







					"""simple docstring"""






					A__     =  TimesformerForVideoClassification.from_pretrained("""facebook/timesformer-base-finetuned-k400"""  ).to(
					    __lowerCAmelCase  )

					A__     =  self.default_image_processor
					A__     =  prepare_video()
					A__     =  image_processor(video[:8]				,      return_tensors="""pt"""  ).to(__lowerCAmelCase  )

					# forward pass
					with torch.no_grad():
							A__     =  model(**__lowerCAmelCase  )

					# verify the logits
					A__     =  torch.Size((1, 4_00)  )
					self.assertEqual(outputs.logits.shape				,      __lowerCAmelCase  )

					A__     =  torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]  ).to(__lowerCAmelCase  )

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










import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser







def   __lowerCamelCase   (				__a  :Optional[int]       )						->							str:






		"""simple docstring"""

		A__     =  {}
		A__     =  tokenizer(example["""content"""]       ,	truncation=__a       )["""input_ids"""]
		A__     =  len(example["""content"""]       ) / len(output["""input_ids"""]       )
		return output


A       :   Union[str, Any]					= HfArgumentParser(PretokenizationArguments)
A       :   str					= parser.parse_args()
if args.num_workers is None:
					A       :   Any					= multiprocessing.cpu_count()
A       :   int					= AutoTokenizer.from_pretrained(args.tokenizer_dir)

A       :   Optional[int]					= time.time()
A       :   int					= load_dataset(args.dataset_name, split='''train''')
print(F'''Dataset loaded in {time.time()-t_start:.2f}s''')

A       :   Optional[Any]					= time.time()
A       :   int					= ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        '''repo_name''',
        '''path''',
        '''copies''',
        '''size''',
        '''content''',
        '''license''',
        '''hash''',
        '''line_mean''',
        '''line_max''',
        '''alpha_frac''',
        '''autogenerated''',
    ],
)
print(F'''Dataset tokenized in {time.time()-t_start:.2f}s''')

A       :   Tuple					= time.time()
ds.push_to_hub(args.tokenized_data_repo)






print(F'''Data pushed to the hub in {time.time()-t_start:.2f}s''')



