"""simple docstring"""





import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
       import sqlitea

       import sqlalchemy


class   UpperCAmelCase_			(					_a):


      def __init__(     self	,    a	,    a	,    a = None	,    a = None	,    a = False	,    **a	,    ) ->      Tuple:
             super().__init__(features=a	,    cache_dir=a	,    keep_in_memory=a	,    **a							)
             lowercase__    :	str     					=			Sql(
                 cache_dir=a	,    features=a	,    sql=a	,    con=a	,    **a	,    )


      def 	_UpperCAmelCase (     self							) ->      Any:
             lowercase__    :	Optional[Any]     					=			None
             lowercase__    :	Optional[Any]     					=			None
             lowercase__    :	Optional[int]     					=			None
             lowercase__    :	Optional[Any]     					=			None

             self.builder.download_and_prepare(
                 download_config=a	,    download_mode=a	,    verification_mode=a	,    base_path=a	,    )

             # Build dataset for splits
             lowercase__    :	Union[str, Any]     					=			self.builder.as_dataset(
                 split='train'	,    verification_mode=a	,    in_memory=self.keep_in_memory							)
             return dataset





class   UpperCAmelCase_			:


      def __init__(     self	,    a	,    a	,    a	,    a = None	,    a = None	,    **a	,    ) ->      Dict:
             if num_proc is not None and num_proc <= 0:
                    raise ValueError(f"""num_proc {num_proc} must be an integer > 0."""							)

             lowercase__    :	Optional[int]     					=			dataset
             lowercase__    :	Any     					=			name
             lowercase__    :	Tuple     					=			con
             lowercase__    :	Optional[int]     					=			batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
             lowercase__    :	Union[str, Any]     					=			num_proc
             lowercase__    :	List[str]     					=			to_sql_kwargs


      def 	_UpperCAmelCase (     self							) ->      int:
             lowercase__    :	Tuple     					=			self.to_sql_kwargs.pop('sql'	,    a							)
             lowercase__    :	List[str]     					=			self.to_sql_kwargs.pop('con'	,    a							)
             lowercase__    :	int     					=			self.to_sql_kwargs.pop('index'	,    a							)

             lowercase__    :	str     					=			self._write(index=a	,    **self.to_sql_kwargs							)
             return written


      def 	_UpperCAmelCase (     self	,    a							) ->      List[str]:
             lowercase__					,   lowercase__					,   lowercase__    :	List[str]     					=			args
             lowercase__    :	List[Any]     					=			{**to_sql_kwargs, 'if_exists': 'append'} if offset > 0 else to_sql_kwargs
             lowercase__    :	str     					=			query_table(
                 table=self.dataset.data	,    key=slice(a	,    offset + self.batch_size							)	,    indices=self.dataset._indices	,    )
             lowercase__    :	Optional[Any]     					=			batch.to_pandas()
             lowercase__    :	Optional[Any]     					=			df.to_sql(self.name	,    self.con	,    index=a	,    **a							)
             return num_rows or len(a							)


      def 	_UpperCAmelCase (     self	,    a	,    **a							) ->      int:
             lowercase__    :	Union[str, Any]     					=			0

             if self.num_proc is None or self.num_proc == 1:
                    for offset in logging.tqdm(
                        range(0	,    len(self.dataset							)	,    self.batch_size							)	,    unit='ba'	,    disable=not logging.is_progress_bar_enabled()	,    desc='Creating SQL from Arrow format'	,    ):
                           written += self._batch_sql((offset, index, to_sql_kwargs)							)
             else:
                    lowercase__					,   lowercase__    :	Union[str, Any]     					=			len(self.dataset							), self.batch_size
                    with multiprocessing.Pool(self.num_proc							) as pool:
                           for num_rows in logging.tqdm(
                               pool.imap(
                                   self._batch_sql	,    [(offset, index, to_sql_kwargs) for offset in range(0	,    a	,    a							)]	,    )	,    total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size	,    unit='ba'	,    disable=not logging.is_progress_bar_enabled()	,    desc='Creating SQL from Arrow format'	,    ):
                                  written += num_rows

             return written


"""simple docstring"""





# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import importlib.metadata
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, Union

from packaging import version

from ..utils import is_torch_available, logging


if is_torch_available():
       import torch


_UpperCamelCase     :		int											=					logging.get_logger(__name__)


@dataclass
class   UpperCAmelCase_			:


      def __init__(     self	,    a=False	,    a=False	,    a=6.0	,    a=None	,    a=False	,    a=False	,    a=None	,    a="fp4"	,    a=False	,    **a	,    ) ->      Tuple:
             lowercase__    :	str     					=			load_in_abit
             lowercase__    :	str     					=			load_in_abit
             lowercase__    :	List[str]     					=			llm_inta_threshold
             lowercase__    :	Dict     					=			llm_inta_skip_modules
             lowercase__    :	Tuple     					=			llm_inta_enable_fpaa_cpu_offload
             lowercase__    :	Any     					=			llm_inta_has_fpaa_weight
             lowercase__    :	Any     					=			bnb_abit_quant_type
             lowercase__    :	Dict     					=			bnb_abit_use_double_quant

             if bnb_abit_compute_dtype is None:
                    lowercase__    :	Dict     					=			torch.floataa
             elif isinstance(a	,    a							):
                    lowercase__    :	Any     					=			getattr(a	,    a							)
             elif isinstance(a	,    torch.dtype							):
                    lowercase__    :	Any     					=			bnb_abit_compute_dtype
             else:
                    raise ValueError('bnb_4bit_compute_dtype must be a string or a torch.dtype'							)

             self.post_init()


      def 	_UpperCAmelCase (     self							) ->      str:
             if not isinstance(self.llm_inta_threshold	,    a							):
                    raise ValueError('llm_int8_threshold must be a float'							)

             if self.llm_inta_skip_modules is not None and not isinstance(self.llm_inta_skip_modules	,    a							):
                    raise ValueError('llm_int8_skip_modules must be a list of strings'							)
             if not isinstance(self.llm_inta_enable_fpaa_cpu_offload	,    a							):
                    raise ValueError('llm_int8_enable_fp32_cpu_offload must be a boolean'							)

             if not isinstance(self.llm_inta_has_fpaa_weight	,    a							):
                    raise ValueError('llm_int8_has_fp16_weight must be a boolean'							)

             if self.bnb_abit_compute_dtype is not None and not isinstance(self.bnb_abit_compute_dtype	,    torch.dtype							):
                    raise ValueError('bnb_4bit_compute_dtype must be torch.dtype'							)

             if not isinstance(self.bnb_abit_quant_type	,    a							):
                    raise ValueError('bnb_4bit_quant_type must be a string'							)

             if not isinstance(self.bnb_abit_use_double_quant	,    a							):
                    raise ValueError('bnb_4bit_use_double_quant must be a boolean'							)

             if self.load_in_abit and not version.parse(importlib.metadata.version('bitsandbytes'							)							) >= version.parse(
                 '0.39.0'							):
                    raise ValueError(
                        '4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version'							)


      def 	_UpperCAmelCase (     self							) ->      Tuple:
             return self.load_in_abit or self.load_in_abit


      def 	_UpperCAmelCase (     self							) ->      List[str]:
             if self.load_in_abit:
                    return "llm_int8"
             elif self.load_in_abit and self.bnb_abit_quant_type == "fp4":
                    return "fp4"
             elif self.load_in_abit and self.bnb_abit_quant_type == "nf4":
                    return "nf4"
             else:
                    return None


      @classmethod
      def 	_UpperCAmelCase (     cls	,    a	,    a	,    **a							) ->      Optional[Any]:
             lowercase__    :	List[Any]     					=			cls(**a							)

             lowercase__    :	Union[str, Any]     					=			[]
             for key, value in kwargs.items():
                    if hasattr(a	,    a							):
                           setattr(a	,    a	,    a							)
                           to_remove.append(a							)
             for key in to_remove:
                    kwargs.pop(a	,    a							)

             if return_unused_kwargs:
                    return config, kwargs
             else:
                    return config


      def 	_UpperCAmelCase (     self	,    a							) ->      Dict:
             with open(a	,    'w'	,    encoding='utf-8'							) as writer:
                    lowercase__    :	Any     					=			self.to_dict()
                    lowercase__    :	str     					=			json.dumps(a	,    indent=2	,    sort_keys=a							) + '\n'

                    writer.write(a							)


      def 	_UpperCAmelCase (     self							) ->      Dict[str, Any]:
             lowercase__    :	Optional[Any]     					=			copy.deepcopy(self.__dict__							)
             lowercase__    :	Any     					=			str(output['bnb_4bit_compute_dtype']							).split('.'							)[1]

             return output


      def __repr__(     self							) ->      Dict:
             return f"""{self.__class__.__name__} {self.to_json_string()}"""


      def 	_UpperCAmelCase (     self	,    a = True							) ->      str:
             if use_diff is True:
                    lowercase__    :	List[Any]     					=			self.to_diff_dict()
             else:
                    lowercase__    :	List[str]     					=			self.to_dict()
             return json.dumps(a	,    indent=2	,    sort_keys=a							) + "\n"


      def 	_UpperCAmelCase (     self							) ->      Dict[str, Any]:
             lowercase__    :	Tuple     					=			self.to_dict()

             # get the default config dict
             lowercase__    :	Optional[Any]     					=			BitsAndBytesConfig().to_dict()

             lowercase__    :	int     					=			{}

             # only serialize values that differ from the default config
             for key, value in config_dict.items():
                    if value != default_config_dict[key]:
                           lowercase__    :	Optional[int]     					=			value

             return serializable_config_dict

