


import json
import os
import tempfile

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features import ArrayaD
from utils import generate_examples, get_duration


__UpperCAmelCase			:			int			      =					(30, 487)
__UpperCAmelCase			:			Optional[int]			      =					(36, 1024)
__UpperCAmelCase			:			Tuple			      =					(100, 100)
__UpperCAmelCase			:			Dict			      =					100

__UpperCAmelCase			:			Dict			      =					datasets.Features(
    {"text": ArrayaD(SHAPE_TEST_1, dtype="float32"), "image": ArrayaD(SHAPE_TEST_2, dtype="float32")}
)

__UpperCAmelCase      ,     __UpperCAmelCase			:			int			      =					os.path.split(__file__)
__UpperCAmelCase			:			Dict			      =					os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		Dict				,      SCREAMING_SNAKE_CASE_   :		List[Any]				,      SCREAMING_SNAKE_CASE_   :		List[Any]							):






       """simple docstring"""





       with ArrowWriter(features=SCREAMING_SNAKE_CASE_				,      path=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)							) as writer:
              for key, record in dummy_data:
                     UpperCamelCase     :						Optional[int] 						=		my_features.encode_example(SCREAMING_SNAKE_CASE_							)
                     writer.write(SCREAMING_SNAKE_CASE_							)
              UpperCamelCase						,			UpperCamelCase     :						Tuple 						=		writer.finalize()


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		Optional[int]				,      SCREAMING_SNAKE_CASE_   :		Tuple							):






       """simple docstring"""





       UpperCamelCase     :						Tuple 						=		datasets.Dataset.from_file(
           filename=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)				,      info=datasets.DatasetInfo(features=SCREAMING_SNAKE_CASE_							)							)
       for _ in dataset:
              pass


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		Optional[Any]				,      SCREAMING_SNAKE_CASE_   :		Tuple							):






       """simple docstring"""





       UpperCamelCase     :						Tuple 						=		datasets.Dataset.from_file(
           filename=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)				,      info=datasets.DatasetInfo(features=SCREAMING_SNAKE_CASE_							)							)
       dataset.set_format('''numpy'''							)
       for _ in dataset:
              pass


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		Tuple				,      SCREAMING_SNAKE_CASE_   :		Optional[Any]							):






       """simple docstring"""





       UpperCamelCase     :						int 						=		1_0
       UpperCamelCase     :						int 						=		datasets.Dataset.from_file(
           filename=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)				,      info=datasets.DatasetInfo(features=SCREAMING_SNAKE_CASE_							)							)
       for i in range(0				,      len(SCREAMING_SNAKE_CASE_							)				,      SCREAMING_SNAKE_CASE_							):
              UpperCamelCase     :						str 						=		dataset[i : i + batch_size]


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		Union[str, Any]				,      SCREAMING_SNAKE_CASE_   :		Any							):






       """simple docstring"""





       UpperCamelCase     :						Tuple 						=		1_0
       UpperCamelCase     :						Dict 						=		datasets.Dataset.from_file(
           filename=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)				,      info=datasets.DatasetInfo(features=SCREAMING_SNAKE_CASE_							)							)
       dataset.set_format('''numpy'''							)
       for i in range(0				,      len(SCREAMING_SNAKE_CASE_							)				,      SCREAMING_SNAKE_CASE_							):
              UpperCamelCase     :						int 						=		dataset[i : i + batch_size]


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		List[Any]				,      SCREAMING_SNAKE_CASE_   :		List[str]							):






       """simple docstring"""





       UpperCamelCase     :						List[str] 						=		datasets.Dataset.from_file(
           filename=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)				,      info=datasets.DatasetInfo(features=SCREAMING_SNAKE_CASE_							)							)
       for col in feats:
              UpperCamelCase     :						Any 						=		dataset[col]


@get_duration
def   a						(				SCREAMING_SNAKE_CASE_   :		Optional[Any]				,      SCREAMING_SNAKE_CASE_   :		str							):






       """simple docstring"""





       UpperCamelCase     :						int 						=		datasets.Dataset.from_file(
           filename=os.path.join(SCREAMING_SNAKE_CASE_				,      '''beta.arrow'''							)				,      info=datasets.DatasetInfo(features=SCREAMING_SNAKE_CASE_							)							)
       dataset.set_format('''numpy'''							)
       for col in feats:
              UpperCamelCase     :						str 						=		dataset[col]


def   a						(				):






       """simple docstring"""





       UpperCamelCase     :						Dict 						=		{}
       UpperCamelCase     :						List[str] 						=		(
           read_unformated,
           read_formatted_as_numpy,
           read_batch_unformated,
           read_batch_formatted_as_numpy,
           read_col_unformated,
           read_col_formatted_as_numpy,
       )
       with tempfile.TemporaryDirectory() as tmp_dir:
              UpperCamelCase     :						Union[str, Any] 						=		datasets.Features({'''image''': ArrayaD(SCREAMING_SNAKE_CASE_				,      dtype='''float32'''							)}							)
              UpperCamelCase     :						Optional[int] 						=		generate_examples(features=SCREAMING_SNAKE_CASE_				,      num_examples=SCREAMING_SNAKE_CASE_							)
              UpperCamelCase     :						List[Any] 						=		write(SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_							)
              for read_func in read_functions:
                     UpperCamelCase     :						Union[str, Any] 						=		read_func(SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_							)

       with tempfile.TemporaryDirectory() as tmp_dir:
              # don't use fixed length for fair comparison
              # feats = datasets.Features(
              #     {"image": datasets.Sequence(datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[1]), SPEED_TEST_SHAPE[0])}
              # )
              UpperCamelCase     :						Dict 						=		datasets.Features({'''image''': datasets.Sequence(datasets.Sequence(datasets.Value('''float32'''							)							)							)}							)
              UpperCamelCase     :						Union[str, Any] 						=		generate_examples(
                  features=SCREAMING_SNAKE_CASE_				,      num_examples=SCREAMING_SNAKE_CASE_				,      seq_shapes={'''image''': SPEED_TEST_SHAPE}							)
              UpperCamelCase     :						Dict 						=		write(SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_							)
              for read_func in read_functions:
                     UpperCamelCase     :						Optional[Any] 						=		read_func(SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_							)

       with tempfile.TemporaryDirectory() as tmp_dir:
              # don't use fixed length for fair comparison
              # feats = datasets.Features(
              #     {"image": datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1])}
              # )
              UpperCamelCase     :						str 						=		datasets.Features({'''image''': datasets.Sequence(datasets.Value('''float32'''							)							)}							)
              UpperCamelCase     :						List[str] 						=		generate_examples(
                  features=SCREAMING_SNAKE_CASE_				,      num_examples=SCREAMING_SNAKE_CASE_				,      seq_shapes={'''image''': [SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1]]}				,      )
              UpperCamelCase     :						Tuple 						=		write(SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_							)
              for read_func in read_functions:
                     UpperCamelCase     :						int 						=		read_func(SCREAMING_SNAKE_CASE_				,      SCREAMING_SNAKE_CASE_							)

       with open(SCREAMING_SNAKE_CASE_				,      '''wb'''							) as f:
              f.write(json.dumps(SCREAMING_SNAKE_CASE_							).encode('''utf-8'''							)							)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_array_xd()





import random
import shutil
import tempfile
import unittest

import numpy as np

from transformers import TaTokenizer, TaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torch
from transformers.utils.import_utils import is_speech_available, is_torch_available


if is_torch_available():
  pass

if is_speech_available():
  from transformers import EncodecFeatureExtractor, MusicgenProcessor


__UpperCAmelCase			:			Union[str, Any]			      =					random.Random()


def   a						(				SCREAMING_SNAKE_CASE_   :		Optional[Any]				,      SCREAMING_SNAKE_CASE_   :		Optional[Any]=1.0				,      SCREAMING_SNAKE_CASE_   :		Optional[Any]=None				,      SCREAMING_SNAKE_CASE_   :		Union[str, Any]=None							):






       """simple docstring"""





       if rng is None:
              UpperCamelCase     :						str 						=		global_rng

       UpperCamelCase     :						Union[str, Any] 						=		[]
       for batch_idx in range(shape[0]							):
              values.append([]							)
              for _ in range(shape[1]							):
                     values[-1].append(rng.random() * scale							)

       return values



@require_torch
@require_sentencepiece
class  UpperCAmelCase_						(				unittest.TestCase):
       '''simple docstring'''






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						int 						=		'''facebook/musicgen-small'''
              UpperCamelCase     :						str 						=		tempfile.mkdtemp()






       def 			_lowercase (  self			,       **__SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              return TaTokenizer.from_pretrained(self.checkpoint			,       **__SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self			,       **__SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              return EncodecFeatureExtractor.from_pretrained(self.checkpoint			,       **__SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              shutil.rmtree(self.tmpdirname    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Union[str, Any] 						=		self.get_tokenizer()
              UpperCamelCase     :						Any 						=		self.get_feature_extractor()

              UpperCamelCase     :						Optional[int] 						=		MusicgenProcessor(tokenizer=__SCREAMING_SNAKE_CASE			,       feature_extractor=__SCREAMING_SNAKE_CASE    )

              processor.save_pretrained(self.tmpdirname    )
              UpperCamelCase     :						int 						=		MusicgenProcessor.from_pretrained(self.tmpdirname    )

              self.assertEqual(processor.tokenizer.get_vocab()			,       tokenizer.get_vocab()    )
              self.assertIsInstance(processor.tokenizer			,       __SCREAMING_SNAKE_CASE    )

              self.assertEqual(processor.feature_extractor.to_json_string()			,       feature_extractor.to_json_string()    )
              self.assertIsInstance(processor.feature_extractor			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						str 						=		MusicgenProcessor(tokenizer=self.get_tokenizer()			,       feature_extractor=self.get_feature_extractor()    )
              processor.save_pretrained(self.tmpdirname    )

              UpperCamelCase     :						Any 						=		self.get_tokenizer(bos_token='''(BOS)'''			,       eos_token='''(EOS)'''    )
              UpperCamelCase     :						Union[str, Any] 						=		self.get_feature_extractor(do_normalize=__SCREAMING_SNAKE_CASE			,       padding_value=1.0    )

              UpperCamelCase     :						Union[str, Any] 						=		MusicgenProcessor.from_pretrained(
                  self.tmpdirname			,       bos_token='''(BOS)'''			,       eos_token='''(EOS)'''			,       do_normalize=__SCREAMING_SNAKE_CASE			,       padding_value=1.0    )

              self.assertEqual(processor.tokenizer.get_vocab()			,       tokenizer_add_kwargs.get_vocab()    )
              self.assertIsInstance(processor.tokenizer			,       __SCREAMING_SNAKE_CASE    )

              self.assertEqual(processor.feature_extractor.to_json_string()			,       feature_extractor_add_kwargs.to_json_string()    )
              self.assertIsInstance(processor.feature_extractor			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						int 						=		self.get_feature_extractor()
              UpperCamelCase     :						Tuple 						=		self.get_tokenizer()

              UpperCamelCase     :						Any 						=		MusicgenProcessor(tokenizer=__SCREAMING_SNAKE_CASE			,       feature_extractor=__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						Union[str, Any] 						=		floats_list((3, 1_000)    )

              UpperCamelCase     :						Tuple 						=		feature_extractor(__SCREAMING_SNAKE_CASE			,       return_tensors='''np'''    )
              UpperCamelCase     :						Optional[Any] 						=		processor(__SCREAMING_SNAKE_CASE			,       return_tensors='''np'''    )

              for key in input_feat_extract.keys():
                     self.assertAlmostEqual(input_feat_extract[key].sum()			,       input_processor[key].sum()			,       delta=1e-2    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						int 						=		self.get_feature_extractor()
              UpperCamelCase     :						List[str] 						=		self.get_tokenizer()

              UpperCamelCase     :						Optional[int] 						=		MusicgenProcessor(tokenizer=__SCREAMING_SNAKE_CASE			,       feature_extractor=__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						Tuple 						=		'''This is a test string'''

              UpperCamelCase     :						Dict 						=		processor(text=__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						Union[str, Any] 						=		tokenizer(__SCREAMING_SNAKE_CASE    )

              for key in encoded_tok.keys():
                     self.assertListEqual(encoded_tok[key]			,       encoded_processor[key]    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Any 						=		self.get_feature_extractor()
              UpperCamelCase     :						int 						=		self.get_tokenizer()

              UpperCamelCase     :						Optional[int] 						=		MusicgenProcessor(tokenizer=__SCREAMING_SNAKE_CASE			,       feature_extractor=__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						int 						=		[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

              UpperCamelCase     :						int 						=		processor.batch_decode(sequences=__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Union[str, Any] 						=		tokenizer.batch_decode(__SCREAMING_SNAKE_CASE    )

              self.assertListEqual(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Dict 						=		self.get_feature_extractor()
              UpperCamelCase     :						Union[str, Any] 						=		self.get_tokenizer()

              UpperCamelCase     :						str 						=		MusicgenProcessor(tokenizer=__SCREAMING_SNAKE_CASE			,       feature_extractor=__SCREAMING_SNAKE_CASE    )

              self.assertListEqual(
                  processor.model_input_names			,       feature_extractor.model_input_names			,       msg='''`processor` and `feature_extractor` model input names do not match'''			,       )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						str 						=		self.get_feature_extractor(padding_side='''left'''    )
              UpperCamelCase     :						Optional[Any] 						=		self.get_tokenizer()

              UpperCamelCase     :						List[Any] 						=		MusicgenProcessor(tokenizer=__SCREAMING_SNAKE_CASE			,       feature_extractor=__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						List[str] 						=		[floats_list((1, x)    )[0] for x in range(5			,       20			,       5    )]
              UpperCamelCase     :						Tuple 						=		processor(__SCREAMING_SNAKE_CASE    ).padding_mask

              UpperCamelCase     :						List[str] 						=		np.asarray(floats_list((3, 20)    )    )[:, None, :]
              UpperCamelCase     :						Optional[int] 						=		processor.batch_decode(__SCREAMING_SNAKE_CASE			,       padding_mask=__SCREAMING_SNAKE_CASE    )

              self.assertIsInstance(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              for audio in decoded_audios:
                     self.assertIsInstance(__SCREAMING_SNAKE_CASE			,       np.ndarray    )

              self.assertTrue(decoded_audios[0].shape == (1, 10)    )
              self.assertTrue(decoded_audios[1].shape == (1, 15)    )
              self.assertTrue(decoded_audios[2].shape == (1, 20)    )

