

'''simple docstring'''







from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


__lowerCamelCase						    =      logging.get_logger(__name__)
def  UpperCAmelCase__      ( UpperCAmelCase__,    UpperCAmelCase__			)     ->     Union[str, Any]:







       try:
              with open(UpperCamelCase__,    """rb"""			) as flax_state_f:
                     A_		   = from_bytes(UpperCamelCase__,    flax_state_f.read()			)
       except UnpicklingError as e:
              try:
                     with open(UpperCamelCase__			) as f:
                            if f.read().startswith("""version"""			):
                                   raise OSError(
                                       """You seem to have cloned a repository without having git-lfs installed. Please"""
                                       """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
                                       """ folder you cloned."""			)
                            else:
                                   raise ValueError from e
              except (UnicodeDecodeError, ValueError):
                     raise EnvironmentError(F'''Unable to convert {model_file} to Flax deserializable object. '''			)

       return load_flax_weights_in_pytorch_model(UpperCamelCase__,    UpperCamelCase__			)







def  UpperCAmelCase__      ( UpperCAmelCase__,    UpperCAmelCase__			)     ->     List[Any]:







       try:
              import torch  # noqa: F401
       except ImportError:
              logger.error(
                  """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
                  """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
                  """ instructions."""			)
              raise

       # check if we have bf16 weights
       A_		   = flatten_dict(jax.tree_util.tree_map(lambda UpperCAmelCase__			: x.dtype == jnp.bfloataa,    UpperCamelCase__			)			).values()
       if any(UpperCamelCase__			):
              # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

              # and bf16 is not fully supported in PT yet.
              logger.warning(
                  """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
                  """before loading those in PyTorch model."""			)
              A_		   = jax.tree_util.tree_map(
                  lambda UpperCAmelCase__			: params.astype(np.floataa			) if params.dtype == jnp.bfloataa else params,    UpperCamelCase__			)

       A_		   = """"""

       A_		   = flatten_dict(UpperCamelCase__,    sep="""."""			)
       A_		   = pt_model.state_dict()

       # keep track of unexpected & missing keys
       A_		   = []
       A_		   = set(pt_model_dict.keys()			)

       for flax_key_tuple, flax_tensor in flax_state_dict.items():
              A_		   = flax_key_tuple.split("""."""			)

              if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
                     A_		   = flax_key_tuple_array[:-1] + ["""weight"""]
                     A_		   = jnp.transpose(UpperCamelCase__,    (3, 2, 0, 1)			)
              elif flax_key_tuple_array[-1] == "kernel":
                     A_		   = flax_key_tuple_array[:-1] + ["""weight"""]
                     A_		   = flax_tensor.T
              elif flax_key_tuple_array[-1] == "scale":
                     A_		   = flax_key_tuple_array[:-1] + ["""weight"""]

              if "time_embedding" not in flax_key_tuple_array:
                     for i, flax_key_tuple_string in enumerate(UpperCamelCase__			):
                            A_		   = (
                                flax_key_tuple_string.replace("""_0""",    """.0"""			)
                                .replace("""_1""",    """.1"""			)
                                .replace("""_2""",    """.2"""			)
                                .replace("""_3""",    """.3"""			)
                                .replace("""_4""",    """.4"""			)
                                .replace("""_5""",    """.5"""			)
                                .replace("""_6""",    """.6"""			)
                                .replace("""_7""",    """.7"""			)
                                .replace("""_8""",    """.8"""			)
                                .replace("""_9""",    """.9"""			)
                            )

              A_		   = """.""".join(UpperCamelCase__			)

              if flax_key in pt_model_dict:
                     if flax_tensor.shape != pt_model_dict[flax_key].shape:
                            raise ValueError(
                                F'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
                                F'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'''			)
                     else:
                            # add weight to pytorch dict
                            A_		   = np.asarray(UpperCamelCase__			) if not isinstance(UpperCamelCase__,    np.ndarray			) else flax_tensor
                            A_		   = torch.from_numpy(UpperCamelCase__			)
                            # remove from missing keys
                            missing_keys.remove(UpperCamelCase__			)
              else:
                     # weight is not expected by PyTorch model
                     unexpected_keys.append(UpperCamelCase__			)

       pt_model.load_state_dict(UpperCamelCase__			)

       # re-transform missing_keys to list
       A_		   = list(UpperCamelCase__			)

       if len(UpperCamelCase__			) > 0:
              logger.warning(
                  """Some weights of the Flax model were not used when initializing the PyTorch model"""
                  F''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
                  F''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
                  """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
                  F''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
                  """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
                  """ FlaxBertForSequenceClassification model)."""			)
       if len(UpperCamelCase__			) > 0:
              logger.warning(
                  F'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
                  F''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
                  """ use it for predictions and inference."""			)

       return pt_model



'''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    A__ ( _snake_case     ):



    def __init__(   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__ = None						,			UpperCamelCase__ = None						,			UpperCamelCase__ = False						,			**UpperCamelCase__						,			)  -> Optional[Any]:





           '''simple docstring'''
           super().__init__(features=UpperCamelCase__						,			cache_dir=UpperCamelCase__						,			keep_in_memory=UpperCamelCase__						,			**UpperCamelCase__    )
           A_		   = Sql(
               cache_dir=UpperCamelCase__						,			features=UpperCamelCase__						,			sql=UpperCamelCase__						,			con=UpperCamelCase__						,			**UpperCamelCase__						,			)







    def    snake_case_     (   self    )  -> Optional[Any]:





           '''simple docstring'''
           A_		   = None
           A_		   = None
           A_		   = None
           A_		   = None

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

           # Build dataset for splits
           A_		   = self.builder.as_dataset(
               split="""train"""						,			verification_mode=UpperCamelCase__						,			in_memory=self.keep_in_memory    )
           return dataset


class    A__ :



    def __init__(   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__ = None						,			UpperCamelCase__ = None						,			**UpperCamelCase__						,			)  -> List[Any]:





           '''simple docstring'''
           if num_proc is not None and num_proc <= 0:
                  raise ValueError(f'''num_proc {num_proc} must be an integer > 0.'''    )

           A_		   = dataset
           A_		   = name
           A_		   = con
           A_		   = batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
           A_		   = num_proc
           A_		   = to_sql_kwargs



    def    snake_case_     (   self    )  -> int:





           '''simple docstring'''
           A_		   = self.to_sql_kwargs.pop("""sql"""						,			UpperCamelCase__    )
           A_		   = self.to_sql_kwargs.pop("""con"""						,			UpperCamelCase__    )
           A_		   = self.to_sql_kwargs.pop("""index"""						,			UpperCamelCase__    )

           A_		   = self._write(index=UpperCamelCase__						,			**self.to_sql_kwargs    )
           return written



    def    snake_case_     (   self						,			UpperCamelCase__    )  -> Union[str, Any]:





           '''simple docstring'''
           A_   ,     A_   ,     A_		   = args
           A_		   = {**to_sql_kwargs, """if_exists""": """append"""} if offset > 0 else to_sql_kwargs
           A_		   = query_table(
               table=self.dataset.data						,			key=slice(UpperCamelCase__						,			offset + self.batch_size    )						,			indices=self.dataset._indices						,			)
           A_		   = batch.to_pandas()
           A_		   = df.to_sql(self.name						,			self.con						,			index=UpperCamelCase__						,			**UpperCamelCase__    )
           return num_rows or len(UpperCamelCase__    )







    def    snake_case_     (   self						,			UpperCamelCase__						,			**UpperCamelCase__    )  -> int:





           '''simple docstring'''
           A_		   = 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:
                  A_   ,     A_		   = 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						,			UpperCamelCase__						,			UpperCamelCase__    )]						,			)						,			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
