






import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase			) ->	Any:
     if model.config.model_type == "gpt2":
          return model.transformer.h[0].mlp.c_fc
     return model.transformer.h[0].mlp.dense_ah_to_h


if is_torch_available():
 import torch
 import torch.nn as nn





 class 			__a  (    nn.Module       ):







  def __init__(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   nn.Module							,   SCREAMING_SNAKE_CASE					:   int						):


       '''simple docstring'''






       super().__init__()
       UpperCamelCase__			:   List[Any]       				=  module
       UpperCamelCase__			:   Union[str, Any]       				=  nn.Sequential(
           nn.Linear(module.in_features							,   SCREAMING_SNAKE_CASE							,   bias=SCREAMING_SNAKE_CASE						)							,   nn.Linear(SCREAMING_SNAKE_CASE							,   module.out_features							,   bias=SCREAMING_SNAKE_CASE						)							,   )
       UpperCamelCase__			:   List[Any]       				=  (2.0 / (5 * min(module.in_features							,   module.out_features						))) ** 0.5
       nn.init.normal_(self.adapter[0].weight							,   std=SCREAMING_SNAKE_CASE						)
       nn.init.zeros_(self.adapter[1].weight						)
       self.adapter.to(module.weight.device						)







  def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Dict							,   *SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   **SCREAMING_SNAKE_CASE					:   Optional[Any]						):


       '''simple docstring'''






       return self.module(SCREAMING_SNAKE_CASE							,   *SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						) + self.adapter(SCREAMING_SNAKE_CASE						)





@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__a  (    unittest.TestCase       ):
 # We keep the constants inside the init function and model loading inside setUp function

 # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
 # Therefore here we use only bloom-1b3 to test our module
 _lowerCAmelCase       :       Union[str, Any]        = '''bigscience/bloom-1b7'''

 # Constant values
 _lowerCAmelCase       :       Tuple        = 2.1_0_9_6_5_9_5_5_2_6_9_2_5_7_4

 _lowerCAmelCase       :       List[str]        = '''Hello my name is'''
 _lowerCAmelCase       :       int        = set()
 EXPECTED_OUTPUTS.add('''Hello my name is John and I am a professional photographer. I'''       )
 EXPECTED_OUTPUTS.add('''Hello my name is John.\nI am a friend of your father.\n'''       )
 EXPECTED_OUTPUTS.add('''Hello my name is John Doe, I am a student at the University'''       )
 _lowerCAmelCase       :       Dict        = 1_0







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  AutoTokenizer.from_pretrained(self.model_name						)







class 			__a  (    A__       ):







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      super().setUp()

      # Models and tokenizer
      UpperCamelCase__			:   Tuple       				=  AutoModelForCausalLM.from_pretrained(
          self.model_name							,   torch_dtype=torch.floataa							,   device_map="auto"						)
      UpperCamelCase__			:   List[str]       				=  AutoModelForCausalLM.from_pretrained(self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      del self.model_fpaa
      del self.model_abit

      gc.collect()
      torch.cuda.empty_cache()







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Any       				=  self.model_abit.config

      self.assertTrue(hasattr(SCREAMING_SNAKE_CASE							,   "quantization_config"						)						)

      UpperCamelCase__			:   int       				=  config.to_dict()
      UpperCamelCase__			:   Union[str, Any]       				=  config.to_diff_dict()

      UpperCamelCase__			:   List[Any]       				=  config.to_json_string()







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      from bitsandbytes.nn import Paramsabit

      UpperCamelCase__			:   Optional[Any]       				=  self.model_fpaa.get_memory_footprint()
      UpperCamelCase__			:   Optional[int]       				=  self.model_abit.get_memory_footprint()

      self.assertAlmostEqual(mem_fpaa / mem_abit							,   self.EXPECTED_RELATIVE_DIFFERENCE						)
      UpperCamelCase__			:   Union[str, Any]       				=  get_some_linear_layer(self.model_abit						)
      self.assertTrue(linear.weight.__class__ == Paramsabit						)







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      from transformers import TaPreTrainedModel

      self.model_fpaa.get_memory_footprint()
      self.model_abit.get_memory_footprint()

      for name, module in self.model_abit.named_modules():
           if isinstance(SCREAMING_SNAKE_CASE							,   torch.nn.Linear						):
                if name not in ["lm_head"] + TaPreTrainedModel._keep_in_fpaa_modules:
                     # 4-bit parameters are packed in uint8 variables
                     self.assertTrue(module.weight.dtype == torch.uinta						)







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						)
      UpperCamelCase__			:   Dict       				=  self.model_abit.generate(input_ids=encoded_input["input_ids"].to(0						)							,   max_new_tokens=10						)

      self.assertIn(self.tokenizer.decode(output_sequences[0]							,   skip_special_tokens=SCREAMING_SNAKE_CASE						)							,   self.EXPECTED_OUTPUTS						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  BitsAndBytesConfig()
      UpperCamelCase__			:   Tuple       				=  True

      UpperCamelCase__			:   Optional[Any]       				=  AutoModelForCausalLM.from_pretrained(
          self.model_name							,   quantization_config=SCREAMING_SNAKE_CASE							,   device_map="auto"						)

      UpperCamelCase__			:   str       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						)
      UpperCamelCase__			:   Tuple       				=  model_abit_from_config.generate(
          input_ids=encoded_input["input_ids"].to(0						)							,   max_new_tokens=10						)

      self.assertIn(self.tokenizer.decode(output_sequences[0]							,   skip_special_tokens=SCREAMING_SNAKE_CASE						)							,   self.EXPECTED_OUTPUTS						)







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      with self.assertRaises(SCREAMING_SNAKE_CASE						), tempfile.TemporaryDirectory() as tmpdirname:
           self.model_abit.save_pretrained(SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Any       				=  BitsAndBytesConfig()

      with self.assertRaises(SCREAMING_SNAKE_CASE						):
           UpperCamelCase__			:   int       				=  AutoModelForCausalLM.from_pretrained(
               self.model_name							,   quantization_config=SCREAMING_SNAKE_CASE							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"							,   bnb_abit_quant_type="nf4"							,   )







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      with self.assertRaises(SCREAMING_SNAKE_CASE						):
           # Tries with `str`
           self.model_abit.to("cpu"						)

      with self.assertRaises(SCREAMING_SNAKE_CASE						):
           # Tries with a `dtype``
           self.model_abit.to(torch.floataa						)

      with self.assertRaises(SCREAMING_SNAKE_CASE						):
           # Tries with a `device`
           self.model_abit.to(torch.device("cuda:0"						)						)

      with self.assertRaises(SCREAMING_SNAKE_CASE						):
           # Tries with a `device`
           self.model_abit.float()

      with self.assertRaises(SCREAMING_SNAKE_CASE						):
           # Tries with a `device`
           self.model_abit.half()

      # Test if we did not break anything
      UpperCamelCase__			:   Any       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						)

      UpperCamelCase__			:   int       				=  self.model_fpaa.to(torch.floataa						)
      UpperCamelCase__			:   Optional[int]       				=  self.model_fpaa.generate(input_ids=encoded_input["input_ids"].to(0						)							,   max_new_tokens=10						)

      # Check this does not throw an error
      UpperCamelCase__			:   Optional[int]       				=  self.model_fpaa.to("cpu"						)

      # Check this does not throw an error
      UpperCamelCase__			:   List[Any]       				=  self.model_fpaa.half()

      # Check this does not throw an error
      UpperCamelCase__			:   Optional[Any]       				=  self.model_fpaa.float()







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  AutoModelForSeqaSeqLM.from_pretrained("t5-small"							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.floataa						)







@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__a  (    unittest.TestCase       ):







 @classmethod
 def 					__lowercase					(		cls					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  "t5-small"
      UpperCamelCase__			:   Optional[Any]       				=  "google/flan-t5-small"  # flan-t5 uses dense-act instead of dense-relu-dense
      UpperCamelCase__			:   Union[str, Any]       				=  AutoTokenizer.from_pretrained(cls.model_name						)
      UpperCamelCase__			:   Any       				=  "Translate in German: Hello, my dog is cute"







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      gc.collect()
      torch.cuda.empty_cache()







 def 					__lowercase					(		self					:   Dict						):


      '''simple docstring'''






      from transformers import TaForConditionalGeneration

      UpperCamelCase__			:   int       				=  TaForConditionalGeneration._keep_in_fpaa_modules
      UpperCamelCase__			:   int       				=  None

      # test with `t5-small`
      UpperCamelCase__			:   List[Any]       				=  TaForConditionalGeneration.from_pretrained(self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      UpperCamelCase__			:   int       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						).to(0						)
      UpperCamelCase__			:   Any       				=  model.generate(**SCREAMING_SNAKE_CASE						)

      # test with `flan-t5-small`
      UpperCamelCase__			:   str       				=  TaForConditionalGeneration.from_pretrained(
          self.dense_act_model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      UpperCamelCase__			:   int       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						).to(0						)
      UpperCamelCase__			:   Union[str, Any]       				=  model.generate(**SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Tuple       				=  modules







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      import bitsandbytes as bnb

      from transformers import TaForConditionalGeneration

      # test with `t5-small`
      UpperCamelCase__			:   List[str]       				=  TaForConditionalGeneration.from_pretrained(self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)

      # there was a bug with decoders - this test checks that it is fixed
      self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q							,   bnb.nn.Linearabit						)						)

      UpperCamelCase__			:   Union[str, Any]       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						).to(0						)
      UpperCamelCase__			:   Optional[int]       				=  model.generate(**SCREAMING_SNAKE_CASE						)

      # test with `flan-t5-small`
      UpperCamelCase__			:   Union[str, Any]       				=  TaForConditionalGeneration.from_pretrained(
          self.dense_act_model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      UpperCamelCase__			:   Tuple       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						).to(0						)
      UpperCamelCase__			:   str       				=  model.generate(**SCREAMING_SNAKE_CASE						)







class 			__a  (    A__       ):







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      super().setUp()
      # model_name
      UpperCamelCase__			:   Any       				=  "bigscience/bloom-560m"
      UpperCamelCase__			:   List[str]       				=  "t5-small"

      # Different types of model

      UpperCamelCase__			:   List[Any]       				=  AutoModel.from_pretrained(self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      # Sequence classification model
      UpperCamelCase__			:   Tuple       				=  AutoModelForSequenceClassification.from_pretrained(
          self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      # CausalLM model
      UpperCamelCase__			:   List[str]       				=  AutoModelForCausalLM.from_pretrained(self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)
      # Seq2seq model
      UpperCamelCase__			:   Dict       				=  AutoModelForSeqaSeqLM.from_pretrained(
          self.seq_to_seq_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="auto"						)







 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      del self.base_model
      del self.sequence_model
      del self.model_abit
      del self.seq_to_seq_model

      gc.collect()
      torch.cuda.empty_cache()







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      from bitsandbytes.nn import Paramsabit

      self.assertTrue(self.base_model.h[-1].mlp.dense_ah_to_h.weight.__class__ == Paramsabit						)

      # Other heads should be nn.Parameter
      self.assertTrue(self.model_abit.lm_head.weight.__class__ == torch.nn.Parameter						)
      self.assertTrue(self.sequence_model.score.weight.__class__ == torch.nn.Parameter						)
      self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter						)







class 			__a  (    A__       ):







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      super().setUp()







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      del self.pipe

      gc.collect()
      torch.cuda.empty_cache()







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  pipeline(
          "text-generation"							,   model=self.model_name							,   model_kwargs={"device_map": "auto", "load_in_4bit": True, "torch_dtype": torch.floataa}							,   max_new_tokens=self.MAX_NEW_TOKENS							,   )

      # Real second forward pass
      UpperCamelCase__			:   int       				=  self.pipe(self.input_text						)
      self.assertIn(pipeline_output[0]["generated_text"]							,   self.EXPECTED_OUTPUTS						)







@require_torch_multi_gpu
class 			__a  (    A__       ):







 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      super().setUp()







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Any       				=  AutoModelForCausalLM.from_pretrained(
          self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE							,   device_map="balanced"						)

      # Check correct device map
      self.assertEqual(set(model_parallel.hf_device_map.values()						)							,   {0, 1}						)

      # Check that inference pass works on the model
      UpperCamelCase__			:   Optional[Any]       				=  self.tokenizer(self.input_text							,   return_tensors="pt"						)

      # Second real batch
      UpperCamelCase__			:   Any       				=  model_parallel.generate(input_ids=encoded_input["input_ids"].to(0						)							,   max_new_tokens=10						)
      self.assertIn(self.tokenizer.decode(output_parallel[0]							,   skip_special_tokens=SCREAMING_SNAKE_CASE						)							,   self.EXPECTED_OUTPUTS						)







class 			__a  (    A__       ):







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  "facebook/opt-350m"
      super().setUp()







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      if version.parse(importlib.metadata.version("bitsandbytes"						)						) < version.parse("0.37.0"						):
           return

      # Step 1: freeze all parameters
      UpperCamelCase__			:   str       				=  AutoModelForCausalLM.from_pretrained(self.model_name							,   load_in_abit=SCREAMING_SNAKE_CASE						)

      self.assertEqual(set(model.hf_device_map.values()						)							,   {torch.cuda.current_device()}						)

      for param in model.parameters():
           UpperCamelCase__			:   Optional[int]       				=  False  # freeze the model - train adapters later
           if param.ndim == 1:
                # cast the small parameters (e.g. layernorm) to fp32 for stability
                UpperCamelCase__			:   List[Any]       				=  param.data.to(torch.floataa						)

        # Step 2: add adapters
      for _, module in model.named_modules():
           if "OPTAttention" in repr(type(SCREAMING_SNAKE_CASE						)						):
                UpperCamelCase__			:   int       				=  LoRALayer(module.q_proj							,   rank=16						)
                UpperCamelCase__			:   List[Any]       				=  LoRALayer(module.k_proj							,   rank=16						)
                UpperCamelCase__			:   List[str]       				=  LoRALayer(module.v_proj							,   rank=16						)

        # Step 3: dummy batch
      UpperCamelCase__			:   int       				=  self.tokenizer("Test batch "							,   return_tensors="pt"						).to(0						)

      # Step 4: Check if the gradient is not None
      with torch.cuda.amp.autocast():
           UpperCamelCase__			:   str       				=  model.forward(**SCREAMING_SNAKE_CASE						)
           out.logits.norm().backward()

      for module in model.modules():
           if isinstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						):
                self.assertTrue(module.adapter[1].weight.grad is not None						)
                self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0						)
           elif isinstance(SCREAMING_SNAKE_CASE							,   nn.Embedding						):
                self.assertTrue(module.weight.grad is None						)







class 			__a  (    A__       ):
 _lowerCAmelCase       :       Any        = '''gpt2-xl'''
 _lowerCAmelCase       :       Union[str, Any]        = 3.3_1_9_1_8_5_4_8_5_4_1_5_2_1_8_7







import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam





class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      return datasets.DatasetInfo(
          features=datasets.Features({"content": datasets.Value("string"						)}						)							,   supervised_keys=SCREAMING_SNAKE_CASE							,   )







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,   gen_kwargs={"examples": get_test_dummy_examples()}						)]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      import apache_beam as beam

      return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE						)







class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      return datasets.DatasetInfo(
          features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"						)}						)}						)							,   supervised_keys=SCREAMING_SNAKE_CASE							,   )







 def 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






      return [
          datasets.SplitGenerator(name=datasets.Split.TRAIN							,   gen_kwargs={"examples": get_test_nested_examples()}						)
      ]







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE						)







def 		SCREAMING_SNAKE_CASE							(					) ->	List[str]:
     return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"]			)]







def 		SCREAMING_SNAKE_CASE							(					) ->	Optional[Any]:
     return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["foo", "bar", "foobar"]			)]





class 			__a  (    A__       ):







 @require_beam
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   List[Any]       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train.arrow'						)						)						)
           self.assertDictEqual(builder.info.features							,   datasets.Features({"content": datasets.Value("string"						)}						)						)
           UpperCamelCase__			:   Optional[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           self.assertDictEqual(dset["train"][0]							,   get_test_dummy_examples()[0][1]						)
           self.assertDictEqual(
               dset["train"][expected_num_examples - 1]							,   get_test_dummy_examples()[expected_num_examples - 1][1]						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset







 @require_beam
 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      UpperCamelCase__			:   Any       				=  beam.io.parquetio.WriteToParquet

      UpperCamelCase__			:   Tuple       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   str       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           with patch("apache_beam.io.parquetio.WriteToParquet"						) as write_parquet_mock:
                UpperCamelCase__			:   Optional[Any]       				=  partial(SCREAMING_SNAKE_CASE							,   num_shards=2						)
                builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertDictEqual(builder.info.features							,   datasets.Features({"content": datasets.Value("string"						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           # Order is not preserved when sharding, so we just check that all the elements are there
           self.assertListEqual(sorted(dset["train"]["content"]						)							,   sorted(["foo", "bar", "foobar"]						)						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset







 @require_beam
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   int       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE						)
           self.assertRaises(datasets.builder.MissingBeamOptions							,   builder.download_and_prepare						)







 @require_beam
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  len(get_test_nested_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   Any       				=  NestedBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train.arrow'						)						)						)
           self.assertDictEqual(
               builder.info.features							,   datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"						)}						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_SNAKE_CASE						)
           self.assertDictEqual(dset["train"][0]							,   get_test_nested_examples()[0][1]						)
           self.assertDictEqual(
               dset["train"][expected_num_examples - 1]							,   get_test_nested_examples()[expected_num_examples - 1][1]						)
           self.assertTrue(
               os.path.exists(os.path.join(SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   "dataset_info.json"						)						)						)
           del dset