





import copy
import os
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union


if TYPE_CHECKING:
     from ...processing_utils import ProcessorMixin
     from ...utils import TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A             =       logging.get_logger(__name__)

__A             =       {
    "google/owlvit-base-patch32": "https://huggingface.co/google/owlvit-base-patch32/resolve/main/config.json",
    "google/owlvit-base-patch16": "https://huggingface.co/google/owlvit-base-patch16/resolve/main/config.json",
    "google/owlvit-large-patch14": "https://huggingface.co/google/owlvit-large-patch14/resolve/main/config.json",
}



class      _SCREAMING_SNAKE_CASE		(			__SCREAMING_SNAKE_CASE					):







 '''simple docstring'''
 lowercase_	=    "owlvit_text_model"

 def __init__(self       :     Any ,				UpperCAmelCase_       :     Optional[int]=49_408 ,				UpperCAmelCase_       :     int=512 ,				UpperCAmelCase_       :     List[str]=2_048 ,				UpperCAmelCase_       :     Union[str, Any]=12 ,				UpperCAmelCase_       :     Optional[Any]=8 ,				UpperCAmelCase_       :     Optional[int]=16 ,				UpperCAmelCase_       :     str="quick_gelu" ,				UpperCAmelCase_       :     int=1E-5 ,				UpperCAmelCase_       :     List[Any]=0.0 ,				UpperCAmelCase_       :     List[str]=0.02 ,				UpperCAmelCase_       :     int=1.0 ,				UpperCAmelCase_       :     List[Any]=0 ,				UpperCAmelCase_       :     str=49_406 ,				UpperCAmelCase_       :     Optional[int]=49_407 ,				**UpperCAmelCase_       :     Optional[Any] ,				) ->Tuple:



   '''simple docstring'''



   super().__init__(pad_token_id=UpperCAmelCase_ ,				bos_token_id=UpperCAmelCase_ ,				eos_token_id=UpperCAmelCase_ ,				**UpperCAmelCase_)

   lowerCamelCase__:   Optional[Any]    						=vocab_size
   lowerCamelCase__:   List[Any]    						=hidden_size
   lowerCamelCase__:   int    						=intermediate_size
   lowerCamelCase__:   Tuple    						=num_hidden_layers
   lowerCamelCase__:   List[str]    						=num_attention_heads
   lowerCamelCase__:   List[str]    						=max_position_embeddings
   lowerCamelCase__:   Union[str, Any]    						=hidden_act
   lowerCamelCase__:   Optional[int]    						=layer_norm_eps
   lowerCamelCase__:   Dict    						=attention_dropout
   lowerCamelCase__:   Optional[int]    						=initializer_range
   lowerCamelCase__:   int    						=initializer_factor







 @classmethod
 def 							SCREAMING_SNAKE_CASE_      (cls       :     str ,				UpperCAmelCase_       :     Union[str, os.PathLike] ,				**UpperCAmelCase_       :     int) ->"PretrainedConfig":



   '''simple docstring'''



   cls._set_token_in_kwargs(UpperCAmelCase_)

   lowerCamelCase__:   Union[str, Any]    						=cls.get_config_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)

   # get the text config dict if we are loading from OwlViTConfig
   if config_dict.get("model_type") == "owlvit":
     lowerCamelCase__:   List[str]    						=config_dict["text_config"]

   if "model_type" in config_dict and hasattr(cls ,				"model_type") and config_dict["model_type"] != cls.model_type:
     logger.warning(
         F"""You are using a model of type {config_dict["model_type"]} to instantiate a model of type """
         F"""{cls.model_type}. This is not supported for all configurations of models and can yield errors.""")

   return cls.from_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)



class      _SCREAMING_SNAKE_CASE		(			__SCREAMING_SNAKE_CASE					):







 '''simple docstring'''
 lowercase_	=    "owlvit_vision_model"

 def __init__(self       :     Dict ,				UpperCAmelCase_       :     Any=768 ,				UpperCAmelCase_       :     Optional[Any]=3_072 ,				UpperCAmelCase_       :     Optional[int]=12 ,				UpperCAmelCase_       :     List[str]=12 ,				UpperCAmelCase_       :     Union[str, Any]=3 ,				UpperCAmelCase_       :     List[str]=768 ,				UpperCAmelCase_       :     Optional[Any]=32 ,				UpperCAmelCase_       :     List[Any]="quick_gelu" ,				UpperCAmelCase_       :     str=1E-5 ,				UpperCAmelCase_       :     List[str]=0.0 ,				UpperCAmelCase_       :     int=0.02 ,				UpperCAmelCase_       :     int=1.0 ,				**UpperCAmelCase_       :     Any ,				) ->Optional[int]:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)

   lowerCamelCase__:   Optional[int]    						=hidden_size
   lowerCamelCase__:   Optional[int]    						=intermediate_size
   lowerCamelCase__:   Union[str, Any]    						=num_hidden_layers
   lowerCamelCase__:   Optional[Any]    						=num_attention_heads
   lowerCamelCase__:   List[str]    						=num_channels
   lowerCamelCase__:   str    						=image_size
   lowerCamelCase__:   Any    						=patch_size
   lowerCamelCase__:   Any    						=hidden_act
   lowerCamelCase__:   int    						=layer_norm_eps
   lowerCamelCase__:   int    						=attention_dropout
   lowerCamelCase__:   Tuple    						=initializer_range
   lowerCamelCase__:   Any    						=initializer_factor







 @classmethod
 def 							SCREAMING_SNAKE_CASE_      (cls       :     Optional[int] ,				UpperCAmelCase_       :     Union[str, os.PathLike] ,				**UpperCAmelCase_       :     Union[str, Any]) ->"PretrainedConfig":



   '''simple docstring'''



   cls._set_token_in_kwargs(UpperCAmelCase_)

   lowerCamelCase__:   str    						=cls.get_config_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)

   # get the vision config dict if we are loading from OwlViTConfig
   if config_dict.get("model_type") == "owlvit":
     lowerCamelCase__:   str    						=config_dict["vision_config"]

   if "model_type" in config_dict and hasattr(cls ,				"model_type") and config_dict["model_type"] != cls.model_type:
     logger.warning(
         F"""You are using a model of type {config_dict["model_type"]} to instantiate a model of type """
         F"""{cls.model_type}. This is not supported for all configurations of models and can yield errors.""")

   return cls.from_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)



class      _SCREAMING_SNAKE_CASE		(			__SCREAMING_SNAKE_CASE					):







 '''simple docstring'''
 lowercase_	=    "owlvit"
 lowercase_	=    True

 def __init__(self       :     Tuple ,				UpperCAmelCase_       :     List[Any]=None ,				UpperCAmelCase_       :     List[Any]=None ,				UpperCAmelCase_       :     Optional[int]=512 ,				UpperCAmelCase_       :     Tuple=2.6592 ,				UpperCAmelCase_       :     Union[str, Any]=True ,				**UpperCAmelCase_       :     Dict ,				) ->Optional[int]:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)

   if text_config is None:
     lowerCamelCase__:   List[Any]    						={}
     logger.info("text_config is None. Initializing the OwlViTTextConfig with default values.")

   if vision_config is None:
     lowerCamelCase__:   List[str]    						={}
     logger.info("vision_config is None. initializing the OwlViTVisionConfig with default values.")

   lowerCamelCase__:   List[Any]    						=OwlViTTextConfig(**UpperCAmelCase_)
   lowerCamelCase__:   Dict    						=OwlViTVisionConfig(**UpperCAmelCase_)

   lowerCamelCase__:   Optional[Any]    						=projection_dim
   lowerCamelCase__:   Optional[Any]    						=logit_scale_init_value
   lowerCamelCase__:   List[str]    						=return_dict
   lowerCamelCase__:   Union[str, Any]    						=1.0

 @classmethod
 def 							SCREAMING_SNAKE_CASE_      (cls       :     Optional[Any] ,				UpperCAmelCase_       :     Union[str, os.PathLike] ,				**UpperCAmelCase_       :     List[Any]) ->"PretrainedConfig":



   '''simple docstring'''



   cls._set_token_in_kwargs(UpperCAmelCase_)

   lowerCamelCase__:   List[str]    						=cls.get_config_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)

   if "model_type" in config_dict and hasattr(cls ,				"model_type") and config_dict["model_type"] != cls.model_type:
     logger.warning(
         F"""You are using a model of type {config_dict["model_type"]} to instantiate a model of type """
         F"""{cls.model_type}. This is not supported for all configurations of models and can yield errors.""")

   return cls.from_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)

 @classmethod
 def 							SCREAMING_SNAKE_CASE_      (cls       :     Dict ,				UpperCAmelCase_       :     Dict ,				UpperCAmelCase_       :     Dict ,				**UpperCAmelCase_       :     Optional[int]) ->Any:



   '''simple docstring'''



   lowerCamelCase__:   int    						={}
   lowerCamelCase__:   Any    						=text_config
   lowerCamelCase__:   List[Any]    						=vision_config

   return cls.from_dict(UpperCAmelCase_ ,				**UpperCAmelCase_)







 def 							SCREAMING_SNAKE_CASE_      (self       :     int) ->str:



   '''simple docstring'''



   lowerCamelCase__:   List[Any]    						=copy.deepcopy(self.__dict__)
   lowerCamelCase__:   Optional[int]    						=self.text_config.to_dict()
   lowerCamelCase__:   List[Any]    						=self.vision_config.to_dict()
   lowerCamelCase__:   Tuple    						=self.__class__.model_type
   return output



class      _SCREAMING_SNAKE_CASE		(			__SCREAMING_SNAKE_CASE					):







 '''simple docstring'''

 @property
 def 							SCREAMING_SNAKE_CASE_      (self       :     List[str]) ->Mapping[str, Mapping[int, str]]:



   '''simple docstring'''



   return OrderedDict(
       [
           ("input_ids", {0: "batch", 1: "sequence"}),
           ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
           ("attention_mask", {0: "batch", 1: "sequence"}),
       ])

 @property
 def 							SCREAMING_SNAKE_CASE_      (self       :     Any) ->Mapping[str, Mapping[int, str]]:



   '''simple docstring'''



   return OrderedDict(
       [
           ("logits_per_image", {0: "batch"}),
           ("logits_per_text", {0: "batch"}),
           ("text_embeds", {0: "batch"}),
           ("image_embeds", {0: "batch"}),
       ])

 @property
 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict) ->float:



   '''simple docstring'''



   return 1E-4

 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[Any] ,				UpperCAmelCase_       :     "ProcessorMixin" ,				UpperCAmelCase_       :     int = -1 ,				UpperCAmelCase_       :     int = -1 ,				UpperCAmelCase_       :     Optional["TensorType"] = None ,				) ->Mapping[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:   Optional[int]    						=super().generate_dummy_inputs(
       processor.tokenizer ,				batch_size=UpperCAmelCase_ ,				seq_length=UpperCAmelCase_ ,				framework=UpperCAmelCase_)
   lowerCamelCase__:   Any    						=super().generate_dummy_inputs(
       processor.image_processor ,				batch_size=UpperCAmelCase_ ,				framework=UpperCAmelCase_)
   return {**text_input_dict, **image_input_dict}







 @property
 def 							SCREAMING_SNAKE_CASE_      (self       :     int) ->int:



   '''simple docstring'''



   return 14








import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


__A             =       datasets.utils.logging.get_logger(__name__)



@dataclass
class      _SCREAMING_SNAKE_CASE		(			datasets.BuilderConfig					):







 '''simple docstring'''
 lowercase_	=    1_0000
 lowercase_	=    None
 lowercase_	=    None



class      _SCREAMING_SNAKE_CASE		(			datasets.ArrowBasedBuilder					):







 '''simple docstring'''
 lowercase_	=    ParquetConfig

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->Any:



   '''simple docstring'''



   return datasets.DatasetInfo(features=self.config.features)

 def 							SCREAMING_SNAKE_CASE_      (self       :     Any ,				UpperCAmelCase_       :     Any) ->Any:



   '''simple docstring'''



   if not self.config.data_files:
     raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}""")
   lowerCamelCase__:   Optional[Any]    						=dl_manager.download_and_extract(self.config.data_files)
   if isinstance(UpperCAmelCase_ ,				(str, list, tuple)):
     lowerCamelCase__:   Any    						=data_files
     if isinstance(UpperCAmelCase_ ,				UpperCAmelCase_):
       lowerCamelCase__:   Union[str, Any]    						=[files]
     # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
     lowerCamelCase__:   Optional[int]    						=[dl_manager.iter_files(UpperCAmelCase_) for file in files]
     return [datasets.SplitGenerator(name=datasets.Split.TRAIN ,				gen_kwargs={"files": files})]
   lowerCamelCase__:   int    						=[]
   for split_name, files in data_files.items():
     if isinstance(UpperCAmelCase_ ,				UpperCAmelCase_):
       lowerCamelCase__:   List[Any]    						=[files]
     # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
     lowerCamelCase__:   str    						=[dl_manager.iter_files(UpperCAmelCase_) for file in files]
     # Infer features is they are stoed in the arrow schema
     if self.info.features is None:
       for file in itertools.chain.from_iterable(UpperCAmelCase_):
         with open(UpperCAmelCase_ ,				"rb") as f:
           lowerCamelCase__:   Union[str, Any]    						=datasets.Features.from_arrow_schema(pq.read_schema(UpperCAmelCase_))
         break
     splits.append(datasets.SplitGenerator(name=UpperCAmelCase_ ,				gen_kwargs={"files": files}))
   return splits

 def 							SCREAMING_SNAKE_CASE_      (self       :     Any ,				UpperCAmelCase_       :     pa.Table) ->pa.Table:



   '''simple docstring'''



   if self.info.features is not None:
     # more expensive cast to support nested features with keys in a different order
     # allows str <-> int/float or str to Audio for example
     lowerCamelCase__:   str    						=table_cast(UpperCAmelCase_ ,				self.info.features.arrow_schema)
   return pa_table







 def 							SCREAMING_SNAKE_CASE_      (self       :     str ,				UpperCAmelCase_       :     Optional[Any]) ->int:



   '''simple docstring'''



   lowerCamelCase__:   str    						=self.info.features.arrow_schema if self.info.features is not None else None
   if self.info.features is not None and self.config.columns is not None:
     if sorted(field.name for field in schema) != sorted(self.config.columns):
       raise ValueError(
           F"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'""")
   for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase_)):
     with open(UpperCAmelCase_ ,				"rb") as f:
       lowerCamelCase__:   Optional[Any]    						=pq.ParquetFile(UpperCAmelCase_)
       try:
         for batch_idx, record_batch in enumerate(
             parquet_file.iter_batches(batch_size=self.config.batch_size ,				columns=self.config.columns)):
           lowerCamelCase__:   Optional[Any]    						=pa.Table.from_batches([record_batch])
           # Uncomment for debugging (will print the Arrow table size and elements)
           # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
           # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
           yield F"""{file_idx}_{batch_idx}""", self._cast_table(UpperCAmelCase_)
       except ValueError as e:
         logger.error(F"""Failed to read file '{file}' with error {type(UpperCAmelCase_)}: {e}""")
         raise

