

'''simple docstring'''







import gc
import unittest

from transformers import CTRLConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
          CTRLForSequenceClassification,
          CTRLLMHeadModel,
          CTRLModel,
      )


class    A__ :



    def __init__(   self						,			UpperCamelCase__						,			UpperCamelCase__=14						,			UpperCamelCase__=7						,			UpperCamelCase__=True						,			UpperCamelCase__=True						,			UpperCamelCase__=True						,			UpperCamelCase__=True						,			UpperCamelCase__=True						,			UpperCamelCase__=99						,			UpperCamelCase__=32						,			UpperCamelCase__=5						,			UpperCamelCase__=4						,			UpperCamelCase__=37						,			UpperCamelCase__="gelu"						,			UpperCamelCase__=0.1						,			UpperCamelCase__=0.1						,			UpperCamelCase__=512						,			UpperCamelCase__=16						,			UpperCamelCase__=2						,			UpperCamelCase__=0.02						,			UpperCamelCase__=3						,			UpperCamelCase__=4						,			UpperCamelCase__=None						,			)  -> List[Any]:





           '''simple docstring'''
           A_		   = parent
           A_		   = batch_size
           A_		   = seq_length
           A_		   = is_training
           A_		   = use_token_type_ids
           A_		   = use_input_mask
           A_		   = use_labels
           A_		   = use_mc_token_ids
           A_		   = vocab_size
           A_		   = hidden_size
           A_		   = num_hidden_layers
           A_		   = num_attention_heads
           A_		   = intermediate_size
           A_		   = hidden_act
           A_		   = hidden_dropout_prob
           A_		   = attention_probs_dropout_prob
           A_		   = max_position_embeddings
           A_		   = type_vocab_size
           A_		   = type_sequence_label_size
           A_		   = initializer_range
           A_		   = num_labels
           A_		   = num_choices
           A_		   = scope
           A_		   = self.vocab_size - 1



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





           '''simple docstring'''
           A_		   = ids_tensor([self.batch_size, self.seq_length]						,			self.vocab_size    )

           A_		   = None
           if self.use_input_mask:
                  A_		   = random_attention_mask([self.batch_size, self.seq_length]    )

           A_		   = None
           if self.use_token_type_ids:
                  A_		   = ids_tensor([self.batch_size, self.seq_length]						,			self.type_vocab_size    )

           A_		   = None
           if self.use_mc_token_ids:
                  A_		   = ids_tensor([self.batch_size, self.num_choices]						,			self.seq_length    )

           A_		   = None
           A_		   = None
           A_		   = None
           if self.use_labels:
                  A_		   = ids_tensor([self.batch_size]						,			self.type_sequence_label_size    )
                  A_		   = ids_tensor([self.batch_size, self.seq_length]						,			self.num_labels    )
                  A_		   = ids_tensor([self.batch_size]						,			self.num_choices    )

           A_		   = self.get_config()

           A_		   = ids_tensor([self.num_hidden_layers, self.num_attention_heads]						,			2    )

           return (
               config,
               input_ids,
               input_mask,
               head_mask,
               token_type_ids,
               mc_token_ids,
               sequence_labels,
               token_labels,
               choice_labels,
           )



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





           '''simple docstring'''
           return CTRLConfig(
               vocab_size=self.vocab_size						,			n_embd=self.hidden_size						,			n_layer=self.num_hidden_layers						,			n_head=self.num_attention_heads						,			n_positions=self.max_position_embeddings						,			pad_token_id=self.pad_token_id						,			)



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





           '''simple docstring'''
           A_		   = CTRLModel(config=UpperCamelCase__    )
           model.to(UpperCamelCase__    )
           model.eval()

           model(UpperCamelCase__						,			token_type_ids=UpperCamelCase__						,			head_mask=UpperCamelCase__    )
           model(UpperCamelCase__						,			token_type_ids=UpperCamelCase__    )
           A_		   = model(UpperCamelCase__    )
           self.parent.assertEqual(result.last_hidden_state.shape						,			(self.batch_size, self.seq_length, self.hidden_size)    )
           self.parent.assertEqual(len(result.past_key_values    )						,			config.n_layer    )



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





           '''simple docstring'''
           A_		   = CTRLLMHeadModel(UpperCamelCase__    )
           model.to(UpperCamelCase__    )
           model.eval()

           A_		   = model(UpperCamelCase__						,			token_type_ids=UpperCamelCase__						,			labels=UpperCamelCase__    )
           self.parent.assertEqual(result.loss.shape						,			()    )
           self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)    )



    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           A_		   = self.prepare_config_and_inputs()

           (
               (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     (
               A_
           )   ,     
           )		   = config_and_inputs

           A_		   = {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """head_mask""": head_mask}

           return config, inputs_dict







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





           '''simple docstring'''
           A_		   = self.num_labels
           A_		   = CTRLForSequenceClassification(UpperCamelCase__    )
           model.to(UpperCamelCase__    )
           model.eval()
           A_		   = ids_tensor([self.batch_size]						,			self.type_sequence_label_size    )
           A_		   = model(UpperCamelCase__						,			token_type_ids=UpperCamelCase__						,			labels=UpperCamelCase__    )
           self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)    )


@require_torch
class    A__ ( _snake_case  ,  _snake_case  ,  _snake_case  ,  unittest.TestCase     ):
    lowercase		   =						(CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
    lowercase		   =						(CTRLLMHeadModel,) if is_torch_available() else ()
    lowercase		   =						(
        {
            "feature-extraction": CTRLModel,
            "text-classification": CTRLForSequenceClassification,
            "text-generation": CTRLLMHeadModel,
            "zero-shot": CTRLForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowercase		   =						True
    lowercase		   =						False
    lowercase		   =						False



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





           '''simple docstring'''
           if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
                  # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
                  # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
                  # config could not be created.
                  return True

           return False



    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           A_		   = CTRLModelTester(self    )
           A_		   = ConfigTester(self						,			config_class=UpperCamelCase__						,			n_embd=37    )



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





           '''simple docstring'''
           super().tearDown()
           # clean-up as much as possible GPU memory occupied by PyTorch
           gc.collect()
           torch.cuda.empty_cache()



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





           '''simple docstring'''
           self.config_tester.run_common_tests()



    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_ctrl_model(*UpperCamelCase__    )



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





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_lm_head_model(*UpperCamelCase__    )



    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
    def    snake_case_     (   self    )  -> Union[str, Any]:





           '''simple docstring'''
           pass



    @slow
    def    snake_case_     (   self    )  -> int:





           '''simple docstring'''
           for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_		   = CTRLModel.from_pretrained(UpperCamelCase__    )
                  self.assertIsNotNone(UpperCamelCase__    )







    @unittest.skip("""The model doesn't support left padding"""    )  # and it's not used enough to be worth fixing :)
    def    snake_case_     (   self    )  -> List[Any]:





           '''simple docstring'''
           pass


@require_torch
class    A__ ( unittest.TestCase     ):



    def    snake_case_     (   self    )  -> Any:





           '''simple docstring'''
           super().tearDown()
           # clean-up as much as possible GPU memory occupied by PyTorch
           gc.collect()
           torch.cuda.empty_cache()







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





           '''simple docstring'''
           A_		   = CTRLLMHeadModel.from_pretrained("""ctrl"""    )
           model.to(UpperCamelCase__    )
           A_		   = torch.tensor(
               [[11859, 0, 1611, 8]]						,			dtype=torch.long						,			device=UpperCamelCase__    )  # Legal the president is
           A_		   = [
               11859,
               0,
               1611,
               8,
               5,
               150,
               26449,
               2,
               19,
               348,
               469,
               3,
               2595,
               48,
               20740,
               246533,
               246533,
               19,
               30,
               5,
           ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

           A_		   = model.generate(UpperCamelCase__						,			do_sample=UpperCamelCase__    )
           self.assertListEqual(output_ids[0].tolist()						,			UpperCamelCase__    )



'''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
