


import unittest

from transformers import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING, is_vision_available, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
			from PIL import Image
else:

			class 							a						:




										"""simple docstring"""


										@staticmethod
										def        __snake_case   (				*lowerCamelCase			:    List[str]       ,      **lowerCamelCase			:    Optional[int]							)      ->       Union[str, Any]:
																	pass

@is_pipeline_test
@require_vision
@require_torch
class 							a						(unittest.TestCase ):




							"""simple docstring"""


							__UpperCAmelCase       :				Any	    =    MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

							def        __snake_case   (				self			:    Optional[Any]       ,      lowerCamelCase			:    List[Any]       ,      lowerCamelCase			:    Optional[int]       ,      lowerCamelCase			:    Any							)      ->       Any:
														__snake_case						:     List[str]									=  pipeline(
														    "zero-shot-object-detection"       ,      model="hf-internal-testing/tiny-random-owlvit-object-detection"							)

														__snake_case						:     Dict									=  [
														    {
														        "image": "./tests/fixtures/tests_samples/COCO/000000039769.png",
														        "candidate_labels": ["cat", "remote", "couch"],
														    }
														]
														return object_detector, examples

							def        __snake_case   (				self			:    Tuple       ,      lowerCamelCase			:    Tuple       ,      lowerCamelCase			:    Optional[Any]							)      ->       List[Any]:
														__snake_case						:     Dict									=  object_detector(examples[0]       ,      threshold=0.0							)

														__snake_case						:     Optional[int]									=  len(lowerCamelCase							)
														self.assertGreater(lowerCamelCase       ,      0							)
														self.assertEqual(
														    lowerCamelCase       ,      [
														        {
														            "score": ANY(lowerCamelCase							),
														            "label": ANY(lowerCamelCase							),
														            "box": {"xmin": ANY(lowerCamelCase							), "ymin": ANY(lowerCamelCase							), "xmax": ANY(lowerCamelCase							), "ymax": ANY(lowerCamelCase							)},
														        }
														        for i in range(lowerCamelCase							)
														    ]       ,      )

							@require_tf
							@unittest.skip("Zero Shot Object Detection not implemented in TF"							)
							def        __snake_case   (				self			:    Optional[int]							)      ->       List[str]:
														pass

							@require_torch
							def        __snake_case   (				self			:    Any							)      ->       str:
														__snake_case						:     List[Any]									=  pipeline(
														    "zero-shot-object-detection"       ,      model="hf-internal-testing/tiny-random-owlvit-object-detection"							)

														__snake_case						:     Optional[Any]									=  object_detector(
														    "./tests/fixtures/tests_samples/COCO/000000039769.png"       ,      candidate_labels=["cat", "remote", "couch"]       ,      threshold=0.64       ,      )

														self.assertEqual(
														    nested_simplify(lowerCamelCase       ,      decimals=4							)       ,      [
														        {"score": 0.72_35, "label": "cat", "box": {"xmin": 204, "ymin": 167, "xmax": 232, "ymax": 190}},
														        {"score": 0.72_18, "label": "remote", "box": {"xmin": 204, "ymin": 167, "xmax": 232, "ymax": 190}},
														        {"score": 0.71_84, "label": "couch", "box": {"xmin": 204, "ymin": 167, "xmax": 232, "ymax": 190}},
														        {"score": 0.67_48, "label": "remote", "box": {"xmin": 571, "ymin": 83, "xmax": 598, "ymax": 103}},
														        {"score": 0.66_56, "label": "cat", "box": {"xmin": 571, "ymin": 83, "xmax": 598, "ymax": 103}},
														        {"score": 0.66_14, "label": "couch", "box": {"xmin": 571, "ymin": 83, "xmax": 598, "ymax": 103}},
														        {"score": 0.64_56, "label": "remote", "box": {"xmin": 494, "ymin": 105, "xmax": 521, "ymax": 127}},
														        {"score": 0.6_42, "label": "remote", "box": {"xmin": 67, "ymin": 274, "xmax": 93, "ymax": 297}},
														        {"score": 0.64_19, "label": "cat", "box": {"xmin": 494, "ymin": 105, "xmax": 521, "ymax": 127}},
														    ]       ,      )

														__snake_case						:     str									=  object_detector(
														    [
														        {
														            "image": "./tests/fixtures/tests_samples/COCO/000000039769.png",
														            "candidate_labels": ["cat", "remote", "couch"],
														        }
														    ]       ,      threshold=0.64       ,      )

														self.assertEqual(
														    nested_simplify(lowerCamelCase       ,      decimals=4							)       ,      [
														        [
														            {"score": 0.72_35, "label": "cat", "box": {"xmin": 204, "ymin": 167, "xmax": 232, "ymax": 190}},
														            {"score": 0.72_18, "label": "remote", "box": {"xmin": 204, "ymin": 167, "xmax": 232, "ymax": 190}},
														            {"score": 0.71_84, "label": "couch", "box": {"xmin": 204, "ymin": 167, "xmax": 232, "ymax": 190}},
														            {"score": 0.67_48, "label": "remote", "box": {"xmin": 571, "ymin": 83, "xmax": 598, "ymax": 103}},
														            {"score": 0.66_56, "label": "cat", "box": {"xmin": 571, "ymin": 83, "xmax": 598, "ymax": 103}},
														            {"score": 0.66_14, "label": "couch", "box": {"xmin": 571, "ymin": 83, "xmax": 598, "ymax": 103}},
														            {"score": 0.64_56, "label": "remote", "box": {"xmin": 494, "ymin": 105, "xmax": 521, "ymax": 127}},
														            {"score": 0.6_42, "label": "remote", "box": {"xmin": 67, "ymin": 274, "xmax": 93, "ymax": 297}},
														            {"score": 0.64_19, "label": "cat", "box": {"xmin": 494, "ymin": 105, "xmax": 521, "ymax": 127}},
														        ]
														    ]       ,      )

							@require_torch
							@slow
							def        __snake_case   (				self			:    Union[str, Any]							)      ->       Tuple:
														__snake_case						:     int									=  pipeline("zero-shot-object-detection"							)

														__snake_case						:     Optional[Any]									=  object_detector(
														    "http://images.cocodataset.org/val2017/000000039769.jpg"       ,      candidate_labels=["cat", "remote", "couch"]       ,      )
														self.assertEqual(
														    nested_simplify(lowerCamelCase       ,      decimals=4							)       ,      [
														        {"score": 0.28_68, "label": "cat", "box": {"xmin": 324, "ymin": 20, "xmax": 640, "ymax": 373}},
														        {"score": 0.2_77, "label": "remote", "box": {"xmin": 40, "ymin": 72, "xmax": 177, "ymax": 115}},
														        {"score": 0.25_37, "label": "cat", "box": {"xmin": 1, "ymin": 55, "xmax": 315, "ymax": 472}},
														        {"score": 0.14_74, "label": "remote", "box": {"xmin": 335, "ymin": 74, "xmax": 371, "ymax": 187}},
														        {"score": 0.12_08, "label": "couch", "box": {"xmin": 4, "ymin": 0, "xmax": 642, "ymax": 476}},
														    ]       ,      )

														__snake_case						:     Optional[int]									=  object_detector(
														    [
														        {
														            "image": "http://images.cocodataset.org/val2017/000000039769.jpg",
														            "candidate_labels": ["cat", "remote", "couch"],
														        },
														        {
														            "image": "http://images.cocodataset.org/val2017/000000039769.jpg",
														            "candidate_labels": ["cat", "remote", "couch"],
														        },
														    ]       ,      )
														self.assertEqual(
														    nested_simplify(lowerCamelCase       ,      decimals=4							)       ,      [
														        [
														            {"score": 0.28_68, "label": "cat", "box": {"xmin": 324, "ymin": 20, "xmax": 640, "ymax": 373}},
														            {"score": 0.2_77, "label": "remote", "box": {"xmin": 40, "ymin": 72, "xmax": 177, "ymax": 115}},
														            {"score": 0.25_37, "label": "cat", "box": {"xmin": 1, "ymin": 55, "xmax": 315, "ymax": 472}},
														            {"score": 0.14_74, "label": "remote", "box": {"xmin": 335, "ymin": 74, "xmax": 371, "ymax": 187}},
														            {"score": 0.12_08, "label": "couch", "box": {"xmin": 4, "ymin": 0, "xmax": 642, "ymax": 476}},
														        ],
														        [
														            {"score": 0.28_68, "label": "cat", "box": {"xmin": 324, "ymin": 20, "xmax": 640, "ymax": 373}},
														            {"score": 0.2_77, "label": "remote", "box": {"xmin": 40, "ymin": 72, "xmax": 177, "ymax": 115}},
														            {"score": 0.25_37, "label": "cat", "box": {"xmin": 1, "ymin": 55, "xmax": 315, "ymax": 472}},
														            {"score": 0.14_74, "label": "remote", "box": {"xmin": 335, "ymin": 74, "xmax": 371, "ymax": 187}},
														            {"score": 0.12_08, "label": "couch", "box": {"xmin": 4, "ymin": 0, "xmax": 642, "ymax": 476}},
														        ],
														    ]       ,      )

							@require_tf
							@unittest.skip("Zero Shot Object Detection not implemented in TF"							)
							def        __snake_case   (				self			:    Optional[Any]							)      ->       int:
														pass

							@require_torch
							@slow
							def        __snake_case   (				self			:    Optional[int]							)      ->       Any:
														__snake_case						:     List[Any]									=  0.2
														__snake_case						:     Any									=  pipeline("zero-shot-object-detection"							)

														__snake_case						:     Optional[Any]									=  object_detector(
														    "http://images.cocodataset.org/val2017/000000039769.jpg"       ,      candidate_labels=["cat", "remote", "couch"]       ,      threshold=lowerCamelCase       ,      )
														self.assertEqual(
														    nested_simplify(lowerCamelCase       ,      decimals=4							)       ,      [
														        {"score": 0.28_68, "label": "cat", "box": {"xmin": 324, "ymin": 20, "xmax": 640, "ymax": 373}},
														        {"score": 0.2_77, "label": "remote", "box": {"xmin": 40, "ymin": 72, "xmax": 177, "ymax": 115}},
														        {"score": 0.25_37, "label": "cat", "box": {"xmin": 1, "ymin": 55, "xmax": 315, "ymax": 472}},
														    ]       ,      )






							@require_torch
							@slow
							def        __snake_case   (				self			:    int							)      ->       Union[str, Any]:
														__snake_case						:     str									=  2
														__snake_case						:     str									=  pipeline("zero-shot-object-detection"							)

														__snake_case						:     int									=  object_detector(
														    "http://images.cocodataset.org/val2017/000000039769.jpg"       ,      candidate_labels=["cat", "remote", "couch"]       ,      top_k=lowerCamelCase       ,      )
														self.assertEqual(
														    nested_simplify(lowerCamelCase       ,      decimals=4							)       ,      [
														        {"score": 0.28_68, "label": "cat", "box": {"xmin": 324, "ymin": 20, "xmax": 640, "ymax": 373}},
														        {"score": 0.2_77, "label": "remote", "box": {"xmin": 40, "ymin": 72, "xmax": 177, "ymax": 115}},
														    ]       ,      )





import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


snake_case__				:	Dict					       =					logging.getLogger(__name__)


def    __lowerCamelCase   (      )     ->    Any:
   lowerCamelCase_						:   str		= argparse.ArgumentParser(
       description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""					)
   parser.add_argument(
       """--dataset_name"""						,						type=A__						,						default="""wikitext"""						,						help="""Name of the training. Explore datasets at: hf.co/datasets."""						,						)
   parser.add_argument(
       """--dataset_config"""						,						type=A__						,						default="""wikitext-103-raw-v1"""						,						help="""Configuration name of the dataset."""					)
   parser.add_argument(
       """--tokenizer_name_or_path"""						,						type=A__						,						default="""sayakpaul/unigram-tokenizer-wikitext"""						,						help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""						,						)
   parser.add_argument(
       """--shard_size"""						,						type=A__						,						default=1000						,						help="""Number of entries to go in a single shard."""						,						)
   parser.add_argument("""--split"""						,						type=A__						,						default="""train"""						,						choices=["""train""", """test""", """validation"""]					)
   parser.add_argument(
       """--limit"""						,						default=A__						,						type=A__						,						help="""Limit the number of shards (used for debugging)."""						,						)
   parser.add_argument(
       """--max_length"""						,						type=A__						,						default=512						,						help="""Maximum sequence length. For training on TPUs, it helps to have a maximum"""
       """ sequence length that is a multiple of 8."""						,						)
   parser.add_argument(
       """--output_dir"""						,						default="""tf-tpu"""						,						type=A__						,						help="""Output directory where the TFRecord shards will be saved. If the"""
       """ path is appended with `gs://` ('gs://tf-tpu', for example) then the TFRecord"""
       """ shards will be directly saved to a Google Cloud Storage bucket."""						,						)

   lowerCamelCase_						:   Optional[int]		= parser.parse_args()
   return args


def    __lowerCamelCase   (      A__						: Tuple					)     ->    str:
   def fn(A__						: Dict					):
      return tokenizer(examples["""text"""]					)

   return fn


def    __lowerCamelCase   (      A__						: Union[str, Any]					)     ->    Union[str, Any]:
   lowerCamelCase_						:   Dict		= []
   for i in range(len(tokenized_data["""input_ids"""]					)					):
      lowerCamelCase_						:   List[str]		= {
          """input_ids""": tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data["""input_ids"""][i]					)					),
          """attention_mask""": tf.train.Feature(
              intaa_list=tf.train.IntaaList(value=tokenized_data["""attention_mask"""][i]					)					),
      }
      lowerCamelCase_						:   Optional[int]		= tf.train.Features(feature=A__					)
      lowerCamelCase_						:   Tuple		= tf.train.Example(features=A__					)
      lowerCamelCase_						:   Optional[Any]		= example.SerializeToString()
      records.append(A__					)
   return records


def    __lowerCamelCase   (      A__						: Tuple					)     ->    Union[str, Any]:
   lowerCamelCase_						:   int		= datasets.load_dataset(args.dataset_name						,						args.dataset_config						,						split=args.split					)

   if args.limit is not None:
      lowerCamelCase_						:   str		= min(len(A__					)						,						args.limit					)
      lowerCamelCase_						:   Dict		= dataset.select(range(A__					)					)
      print(f'''Limiting the dataset to {args.limit} entries.'''					)

   lowerCamelCase_						:   str		= AutoTokenizer.from_pretrained(args.tokenizer_name_or_path					)

   # Handle output directory creation.
   # For serializing into a Google Cloud Storage Bucket, one needs to first
   # create a bucket.
   if "gs" not in args.output_dir:
      if not os.path.exists(args.output_dir					):
         os.makedirs(args.output_dir					)
      lowerCamelCase_						:   Optional[int]		= os.path.join(args.output_dir						,						args.split					)
      if not os.path.exists(A__					):
         os.makedirs(A__					)
   else:
      lowerCamelCase_						:   Any		= os.path.join(args.output_dir						,						args.split					)

   # Tokenize the whole dataset at once.
   lowerCamelCase_						:   List[str]		= tokenize_function(A__					)
   lowerCamelCase_						:   List[Any]		= dataset.map(A__						,						batched=A__						,						num_proc=4						,						remove_columns=["""text"""]					)

   # We need to concatenate all our texts together, and then split the result
   # into chunks of a fixed size, which we will call block_size. To do this, we
   # will use the map method again, with the option batched=True. When we use batched=True,
   # the function we pass to map() will be passed multiple inputs at once, allowing us
   # to group them into more or fewer examples than we had in the input.
   # This allows us to create our new fixed-length samples. The advantage of this
   # method is that we don't lose a whole lot of content from the dataset compared to the
   # case where we simply tokenize with a pre-defined max_length.

   def group_texts(A__						: Union[str, Any]					):
      # Concatenate all texts.
      lowerCamelCase_						:   Optional[int]		= {k: sum(examples[k]						,						[]					) for k in examples.keys()}
      lowerCamelCase_						:   Union[str, Any]		= len(concatenated_examples[list(examples.keys()					)[0]]					)
      # We drop the small remainder, though you could add padding instead if the model supports it
      # In this, as in all things, we advise you to follow your heart 🫀
      lowerCamelCase_						:   Union[str, Any]		= (total_length // args.max_length) * args.max_length
      # Split by chunks of max_len.
      lowerCamelCase_						:   List[str]		= {
          k: [t[i : i + args.max_length] for i in range(0						,						A__						,						args.max_length					)]
          for k, t in concatenated_examples.items()
      }
      return result

   lowerCamelCase_						:   Optional[int]		= dataset_tokenized.map(A__						,						batched=A__						,						batch_size=1000						,						num_proc=4					)

   lowerCamelCase_						:   List[str]		= 0
   lowerCamelCase_						:   Tuple		= 0
   for shard in range(0						,						len(A__					)						,						args.shard_size					):
      lowerCamelCase_						:   Union[str, Any]		= grouped_dataset[shard : shard + args.shard_size]
      lowerCamelCase_						:   Optional[int]		= len(dataset_snapshot["""input_ids"""]					)
      lowerCamelCase_						:   List[str]		= os.path.join(A__						,						f'''dataset-{shard_count}-{records_containing}.tfrecord'''					)
      lowerCamelCase_						:   Optional[int]		= get_serialized_examples(A__					)

      with tf.io.TFRecordWriter(A__					) as out_file:
         for i in range(len(A__					)					):
            lowerCamelCase_						:   Dict		= serialized_examples[i]
            out_file.write(A__					)
         print("""Wrote file {} containing {} records""".format(A__						,						A__					)					)

      shard_count += 1
      total_records += records_containing

   with open(f'''split-{args.split}-records-count.txt'''						,						"""w"""					) as f:
      print(f'''Total {args.split} records: {total_records}'''						,						file=A__					)


if __name__ == "__main__":
      snake_case__				:	str					       =					parse_args()
      main(args)





