





'''simple docstring'''




import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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 (
        SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        SqueezeBertForMaskedLM,
        SqueezeBertForMultipleChoice,
        SqueezeBertForQuestionAnswering,
        SqueezeBertForSequenceClassification,
        SqueezeBertForTokenClassification,
        SqueezeBertModel,
    )






class UpperCAmelCase							(		_lowerCamelCase							):



    def __init__(  self     ,			__A     ,			__A=13     ,			__A=7     ,			__A=True     ,			__A=True     ,			__A=False     ,			__A=True     ,			__A=99     ,			__A=32     ,			__A=5     ,			__A=4     ,			__A=64     ,			__A="gelu"     ,			__A=0.1     ,			__A=0.1     ,			__A=512     ,			__A=16     ,			__A=2     ,			__A=0.0_2     ,			__A=3     ,			__A=4     ,			__A=None     ,			__A=2     ,			__A=2     ,			__A=2     ,			__A=2     ,			__A=4     ,			__A=1     ,			):
     __UpperCAmelCase					=   parent
     __UpperCAmelCase					=   batch_size
     __UpperCAmelCase					=   seq_length
     __UpperCAmelCase					=   is_training
     __UpperCAmelCase					=   use_input_mask
     __UpperCAmelCase					=   use_token_type_ids
     __UpperCAmelCase					=   use_labels
     __UpperCAmelCase					=   vocab_size
     __UpperCAmelCase					=   hidden_size
     __UpperCAmelCase					=   num_hidden_layers
     __UpperCAmelCase					=   num_attention_heads
     __UpperCAmelCase					=   intermediate_size
     __UpperCAmelCase					=   hidden_act
     __UpperCAmelCase					=   hidden_dropout_prob
     __UpperCAmelCase					=   attention_probs_dropout_prob
     __UpperCAmelCase					=   max_position_embeddings
     __UpperCAmelCase					=   type_vocab_size
     __UpperCAmelCase					=   type_sequence_label_size
     __UpperCAmelCase					=   initializer_range
     __UpperCAmelCase					=   num_labels
     __UpperCAmelCase					=   num_choices
     __UpperCAmelCase					=   scope
     __UpperCAmelCase					=   q_groups
     __UpperCAmelCase					=   k_groups
     __UpperCAmelCase					=   v_groups
     __UpperCAmelCase					=   post_attention_groups
     __UpperCAmelCase					=   intermediate_groups
     __UpperCAmelCase					=   output_groups



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size					)

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

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

     __UpperCAmelCase					=   self.get_config()

     return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



    def 					__lowerCamelCase						(  self					):
     return SqueezeBertConfig(
         embedding_size=self.hidden_size     ,			vocab_size=self.vocab_size     ,			hidden_size=self.hidden_size     ,			num_hidden_layers=self.num_hidden_layers     ,			num_attention_heads=self.num_attention_heads     ,			intermediate_size=self.intermediate_size     ,			hidden_act=self.hidden_act     ,			attention_probs_dropout_prob=self.hidden_dropout_prob     ,			attention_dropout=self.attention_probs_dropout_prob     ,			max_position_embeddings=self.max_position_embeddings     ,			initializer_range=self.initializer_range     ,			q_groups=self.q_groups     ,			k_groups=self.k_groups     ,			v_groups=self.v_groups     ,			post_attention_groups=self.post_attention_groups     ,			intermediate_groups=self.intermediate_groups     ,			output_groups=self.output_groups     ,			)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   SqueezeBertModel(config=A__					)
     model.to(A__					)
     model.eval()
     __UpperCAmelCase					=   model(A__     ,			A__					)
     __UpperCAmelCase					=   model(A__					)
     self.parent.assertEqual(result.last_hidden_state.shape     ,			(self.batch_size, self.seq_length, self.hidden_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   SqueezeBertForMaskedLM(config=A__					)
     model.to(A__					)
     model.eval()
     __UpperCAmelCase					=   model(A__     ,			attention_mask=A__     ,			labels=A__					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   SqueezeBertForQuestionAnswering(config=A__					)
     model.to(A__					)
     model.eval()
     __UpperCAmelCase					=   model(
         A__     ,			attention_mask=A__     ,			start_positions=A__     ,			end_positions=A__					)
     self.parent.assertEqual(result.start_logits.shape     ,			(self.batch_size, self.seq_length)					)
     self.parent.assertEqual(result.end_logits.shape     ,			(self.batch_size, self.seq_length)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   self.num_labels
     __UpperCAmelCase					=   SqueezeBertForSequenceClassification(A__					)
     model.to(A__					)
     model.eval()
     __UpperCAmelCase					=   model(A__     ,			attention_mask=A__     ,			labels=A__					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_labels)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   self.num_labels
     __UpperCAmelCase					=   SqueezeBertForTokenClassification(config=A__					)
     model.to(A__					)
     model.eval()

     __UpperCAmelCase					=   model(A__     ,			attention_mask=A__     ,			labels=A__					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.num_labels)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   self.num_choices
     __UpperCAmelCase					=   SqueezeBertForMultipleChoice(config=A__					)
     model.to(A__					)
     model.eval()
     __UpperCAmelCase					=   input_ids.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   input_mask.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   model(
         A__     ,			attention_mask=A__     ,			labels=A__     ,			)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_choices)					)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.prepare_config_and_inputs()
     ((__UpperCAmelCase)					,    (__UpperCAmelCase)					,    (__UpperCAmelCase)					,    (__UpperCAmelCase)					,    (__UpperCAmelCase)					,    (__UpperCAmelCase))					=   config_and_inputs
     __UpperCAmelCase					=   {'input_ids': input_ids, 'attention_mask': input_mask}
     return config, inputs_dict






@require_torch
class UpperCAmelCase							(		_lowerCamelCase					,			_lowerCamelCase					,			unittest.TestCase							):
    _A :			List[str]						   =       (
        (
            SqueezeBertModel,
            SqueezeBertForMaskedLM,
            SqueezeBertForMultipleChoice,
            SqueezeBertForQuestionAnswering,
            SqueezeBertForSequenceClassification,
            SqueezeBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    _A :			int						   =       (
        {
            """feature-extraction""": SqueezeBertModel,
            """fill-mask""": SqueezeBertForMaskedLM,
            """question-answering""": SqueezeBertForQuestionAnswering,
            """text-classification""": SqueezeBertForSequenceClassification,
            """token-classification""": SqueezeBertForTokenClassification,
            """zero-shot""": SqueezeBertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    _A :			str						   =       False
    _A :			List[Any]						   =       True
    _A :			int						   =       False



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   SqueezeBertModelTester(self					)
     __UpperCAmelCase					=   ConfigTester(self     ,			config_class=A__     ,			dim=37					)



    def 					__lowerCamelCase						(  self					):
     self.config_tester.run_common_tests()



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_model(*A__					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_masked_lm(*A__					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_question_answering(*A__					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_sequence_classification(*A__					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_token_classification(*A__					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_squeezebert_for_multiple_choice(*A__					)




    @slow
    def 					__lowerCamelCase						(  self					):
     for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase					=   SqueezeBertModel.from_pretrained(A__					)
      self.assertIsNotNone(A__					)






@require_sentencepiece
@require_tokenizers
@require_torch
class UpperCAmelCase							(		unittest.TestCase							):



    @slow
    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   SqueezeBertForSequenceClassification.from_pretrained('squeezebert/squeezebert-mnli'					)

     __UpperCAmelCase					=   torch.tensor([[1, 29_414, 232, 328, 740, 1_140, 12_695, 69, 13, 1_588, 2]]					)
     __UpperCAmelCase					=   model(A__					)[0]
     __UpperCAmelCase					=   torch.Size((1, 3)					)
     self.assertEqual(output.shape     ,			A__					)
     __UpperCAmelCase					=   torch.tensor([[0.6_4_0_1, -0.0_3_4_9, -0.6_0_4_1]]					)
     self.assertTrue(torch.allclose(A__     ,			A__     ,			atol=1E-4					)					)












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



    def __init__(  self     ,			__A     ,			__A     ,			__A = None     ,			__A = None     ,			__A = False     ,			**__A     ,			):
     super().__init__(features=__A     ,			cache_dir=__A     ,			keep_in_memory=__A     ,			**__A					)
     __UpperCAmelCase					=   Sql(
         cache_dir=__A     ,			features=__A     ,			sql=__A     ,			con=__A     ,			**__A     ,			)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None

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

     # Build dataset for splits
     __UpperCAmelCase					=   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     ,			):
     if num_proc is not None and num_proc <= 0:
      raise ValueError(f'num_proc {num_proc} must be an integer > 0.'					)

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



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('sql'     ,			__A					)
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('con'     ,			__A					)
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('index'     ,			__A					)

     __UpperCAmelCase					=   self._write(index=__A     ,			**self.to_sql_kwargs					)
     return written



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




    def 					__lowerCamelCase						(  self     ,			__A     ,			**__A					):
     __UpperCAmelCase					=   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:
      __UpperCAmelCase					,    __UpperCAmelCase					=   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





