






"""simple docstring"""


import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


def        UpperCAmelCase      (					a__ ):





						'''simple docstring'''






						lowerCAmelCase   :Union[str, Any]						=			np.inf

						def set_batch_size(a__ ) -> None:
												nonlocal batch_size
												if isinstance(a__     ,      a__ ):
																		lowerCAmelCase   :Any						=			min(a__     ,      config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS )
												elif isinstance(a__     ,      a__ ):
																		lowerCAmelCase   :str						=			min(a__     ,      config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS )
												elif isinstance(a__     ,      a__ ) and feature.dtype == "binary":
																		lowerCAmelCase   :Union[str, Any]						=			min(a__     ,      config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS )

						_visit(a__     ,      a__ )

						return None if batch_size is np.inf else batch_size




class 						__UpperCamelCase     (     UpperCamelCase						):
					def __init__(							self							:	List[str]     , UpperCAmelCase							:	NestedDataStructureLike[PathLike]     , UpperCAmelCase							:	Optional[NamedSplit] = None     , UpperCAmelCase							:	Optional[Features] = None     , UpperCAmelCase							:	str = None     , UpperCAmelCase							:	bool = False     , UpperCAmelCase							:	bool = False     , UpperCAmelCase							:	Optional[int] = None     , **UpperCAmelCase							:	Dict     , )	->       Tuple:
											super().__init__(
											    UpperCAmelCase     , split=UpperCAmelCase     , features=UpperCAmelCase     , cache_dir=UpperCAmelCase     , keep_in_memory=UpperCAmelCase     , streaming=UpperCAmelCase     , num_proc=UpperCAmelCase     , **UpperCAmelCase     , )
											lowerCAmelCase   :List[Any]						=			path_or_paths if isinstance(UpperCAmelCase     , UpperCAmelCase   ) else {self.split: path_or_paths}
											lowerCAmelCase   :List[Any]						=			_PACKAGED_DATASETS_MODULES['parquet'][1]
											lowerCAmelCase   :Dict						=			Parquet(
											    cache_dir=UpperCAmelCase     , data_files=UpperCAmelCase     , features=UpperCAmelCase     , hash=UpperCAmelCase     , **UpperCAmelCase     , )
					def        UpperCAmelCase__ (							self							:	Any   )	->       Optional[Any]:
											# Build iterable dataset
											if self.streaming:
																	lowerCAmelCase   :Optional[Any]						=			self.builder.as_streaming_dataset(split=self.split   )
											# Build regular (map-style) dataset
											else:
																	lowerCAmelCase   :List[Any]						=			None
																	lowerCAmelCase   :List[Any]						=			None
																	lowerCAmelCase   :Any						=			None
																	lowerCAmelCase   :List[str]						=			None

																	self.builder.download_and_prepare(
																	    download_config=UpperCAmelCase     , download_mode=UpperCAmelCase     , verification_mode=UpperCAmelCase     , base_path=UpperCAmelCase     , num_proc=self.num_proc     , )
																	lowerCAmelCase   :Tuple						=			self.builder.as_dataset(
																	    split=self.split     , verification_mode=UpperCAmelCase     , in_memory=self.keep_in_memory   )
											return dataset




class 						__UpperCamelCase     :
					def __init__(							self							:	Optional[int]     , UpperCAmelCase							:	Dataset     , UpperCAmelCase							:	Union[PathLike, BinaryIO]     , UpperCAmelCase							:	Optional[int] = None     , **UpperCAmelCase							:	Optional[Any]     , )	->       Union[str, Any]:
											lowerCAmelCase   :Dict						=			dataset
											lowerCAmelCase   :Union[str, Any]						=			path_or_buf
											lowerCAmelCase   :Optional[int]						=			batch_size or get_writer_batch_size(dataset.features   )
											lowerCAmelCase   :Dict						=			parquet_writer_kwargs
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       int:
											lowerCAmelCase   :Optional[Any]						=			self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

											if isinstance(self.path_or_buf     , (str, bytes, os.PathLike)   ):
																	with open(self.path_or_buf     , 'wb+'   ) as buffer:
																							lowerCAmelCase   :int						=			self._write(file_obj=UpperCAmelCase     , batch_size=UpperCAmelCase     , **self.parquet_writer_kwargs   )
											else:
																	lowerCAmelCase   :Optional[int]						=			self._write(file_obj=self.path_or_buf     , batch_size=UpperCAmelCase     , **self.parquet_writer_kwargs   )
											return written
					def        UpperCAmelCase__ (							self							:	List[Any]     , UpperCAmelCase							:	BinaryIO     , UpperCAmelCase							:	int     , **UpperCAmelCase							:	Union[str, Any]   )	->       int:
											lowerCAmelCase   :List[Any]						=			0
											lowerCAmelCase   :Optional[int]						=			parquet_writer_kwargs.pop('path_or_buf'     , UpperCAmelCase   )
											lowerCAmelCase   :Dict						=			self.dataset.features.arrow_schema

											lowerCAmelCase   :List[str]						=			pq.ParquetWriter(UpperCAmelCase     , schema=UpperCAmelCase     , **UpperCAmelCase   )

											for offset in logging.tqdm(
											    range(0     , len(self.dataset   )     , UpperCAmelCase   )     , unit='ba'     , disable=not logging.is_progress_bar_enabled()     , desc='Creating parquet from Arrow format'     , ):
																	lowerCAmelCase   :List[Any]						=			query_table(
																	    table=self.dataset._data     , key=slice(UpperCAmelCase     , offset + batch_size   )     , indices=self.dataset._indices if self.dataset._indices is not None else None     , )
																	writer.write_table(UpperCAmelCase   )
																	written += batch.nbytes
											writer.close()
											return written







"""simple docstring"""


from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_tf_available():
		from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_VISION_2_SEQ_MAPPING

if is_torch_available():
		import torch

		from ..models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING

__SCREAMING_SNAKE_CASE   =   logging.get_logger(__name__)




@add_end_docstrings(UpperCamelCase						)
class 						__UpperCamelCase     (     UpperCamelCase						):
					def __init__(							self							:	List[Any]     , *UpperCAmelCase							:	Any     , **UpperCAmelCase							:	Tuple   )	->       Optional[int]:
											super().__init__(*UpperCAmelCase     , **UpperCAmelCase   )
											requires_backends(self     , 'vision'   )
											self.check_model_type(
											    TF_MODEL_FOR_VISION_2_SEQ_MAPPING if self.framework == 'tf' else MODEL_FOR_VISION_2_SEQ_MAPPING   )
					def        UpperCAmelCase__ (							self							:	List[str]     , UpperCAmelCase							:	str=None     , UpperCAmelCase							:	str=None     , UpperCAmelCase							:	List[str]=None   )	->       Tuple:
											lowerCAmelCase   :Optional[Any]						=			{}
											lowerCAmelCase   :int						=			{}

											if prompt is not None:
																	lowerCAmelCase   :Any						=			prompt

											if generate_kwargs is not None:
																	lowerCAmelCase   :List[Any]						=			generate_kwargs
											if max_new_tokens is not None:
																	if "generate_kwargs" not in forward_kwargs:
																							lowerCAmelCase   :Any						=			{}
																	if "max_new_tokens" in forward_kwargs["generate_kwargs"]:
																							raise ValueError(
																							    '\'max_new_tokens\' is defined twice, once in \'generate_kwargs\' and once as a direct parameter,'
																							    ' please use only one'   )
																	lowerCAmelCase   :Union[str, Any]						=			max_new_tokens
											return preprocess_params, forward_kwargs, {}
					def __call__(							self							:	Tuple     , UpperCAmelCase							:	Union[str, List[str], "Image.Image", List["Image.Image"]]     , **UpperCAmelCase							:	Any   )	->       Optional[int]:
											return super().__call__(UpperCAmelCase     , **UpperCAmelCase   )
					def        UpperCAmelCase__ (							self							:	List[str]     , UpperCAmelCase							:	Tuple     , UpperCAmelCase							:	Any=None   )	->       str:
											lowerCAmelCase   :Dict						=			load_image(UpperCAmelCase   )

											if prompt is not None:
																	if not isinstance(UpperCAmelCase     , UpperCAmelCase   ):
																							raise ValueError(
																							    f"""Received an invalid text input, got - {type(UpperCAmelCase   )} - but expected a single string. """
																							    'Note also that one single text can be provided for conditional image to text generation.'   )

																	lowerCAmelCase   :List[str]						=			self.model.config.model_type

																	if model_type == "git":
																							lowerCAmelCase   :Any						=			self.image_processor(images=UpperCAmelCase     , return_tensors=self.framework   )
																							lowerCAmelCase   :Union[str, Any]						=			self.tokenizer(text=UpperCAmelCase     , add_special_tokens=UpperCAmelCase   ).input_ids
																							lowerCAmelCase   :Optional[int]						=			[self.tokenizer.cls_token_id] + input_ids
																							lowerCAmelCase   :Union[str, Any]						=			torch.tensor(UpperCAmelCase   ).unsqueeze(0   )
																							model_inputs.update({'input_ids': input_ids}   )

																	elif model_type == "pix2struct":
																							lowerCAmelCase   :Optional[Any]						=			self.image_processor(images=UpperCAmelCase     , header_text=UpperCAmelCase     , return_tensors=self.framework   )

																	elif model_type != "vision-encoder-decoder":
																							# vision-encoder-decoder does not support conditional generation
																							lowerCAmelCase   :int						=			self.image_processor(images=UpperCAmelCase     , return_tensors=self.framework   )
																							lowerCAmelCase   :Optional[Any]						=			self.tokenizer(UpperCAmelCase     , return_tensors=self.framework   )
																							model_inputs.update(UpperCAmelCase   )

																	else:
																							raise ValueError(f"""Model type {model_type} does not support conditional text generation"""   )

											else:
																	lowerCAmelCase   :str						=			self.image_processor(images=UpperCAmelCase     , return_tensors=self.framework   )

											if self.model.config.model_type == "git" and prompt is None:
																	lowerCAmelCase   :Union[str, Any]						=			None

											return model_inputs
					def        UpperCAmelCase__ (							self							:	Optional[Any]     , UpperCAmelCase							:	Union[str, Any]     , UpperCAmelCase							:	Any=None   )	->       int:
											# Git model sets `model_inputs["input_ids"] = None` in `preprocess` (when `prompt=None`). In batch model, the
											# pipeline will group them into a list of `None`, which fail `_forward`. Avoid this by checking it first.
											if (
											    "input_ids" in model_inputs
											    and isinstance(model_inputs['input_ids']     , UpperCAmelCase   )
											    and all(x is None for x in model_inputs['input_ids']   )
											):
																	lowerCAmelCase   :Optional[int]						=			None

											if generate_kwargs is None:
																	lowerCAmelCase   :Optional[int]						=			{}
											# FIXME: We need to pop here due to a difference in how `generation.py` and `generation.tf_utils.py`
											#  parse inputs. In the Tensorflow version, `generate` raises an error if we don't use `input_ids` whereas
											#  the PyTorch version matches it with `self.model.main_input_name` or `self.model.encoder.main_input_name`
											#  in the `_prepare_model_inputs` method.
											lowerCAmelCase   :Optional[int]						=			model_inputs.pop(self.model.main_input_name   )
											lowerCAmelCase   :Any						=			self.model.generate(UpperCAmelCase     , **UpperCAmelCase     , **UpperCAmelCase   )
											return model_outputs
					def        UpperCAmelCase__ (							self							:	Dict     , UpperCAmelCase							:	Union[str, Any]   )	->       Any:
											lowerCAmelCase   :str						=			[]
											for output_ids in model_outputs:
																	lowerCAmelCase   :Optional[int]						=			{
																	    'generated_text': self.tokenizer.decode(
																	        UpperCAmelCase     , skip_special_tokens=UpperCAmelCase     , )
																	}
																	records.append(UpperCAmelCase   )
											return records