





import copy
import os
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, ParquetWriter, TypedSequence
from datasets.features import ArrayaD, ClassLabel, Features, Image, Value
from datasets.features.features import ArrayaDExtensionType, cast_to_python_objects
from datasets.keyhash import DuplicatedKeysError, InvalidKeyError

from .utils import require_pil



class 							_UpperCamelCase					(	lowerCAmelCase_     ):

       def lowercase (		self: List[Any]		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence([1, 2, 3]		)		)
              self.assertEqual(arr.type					,   pa.intaa()		)

       def lowercase (		self: str		)       -> List[str]:




              """simple docstring"""


              with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  pa.array(TypedSequence([1, 2, 3]		)					,   type=pa.intaa()		)

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




              """simple docstring"""


              with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  pa.array(TypedSequence([1, 2, 3]					,   try_type=Value("bool"		)					,   type=Value("int64"		)		)		)

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




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence([1, 2, 3]					,   type=Value("int32"		)		)		)
              self.assertEqual(arr.type					,   pa.intaa()		)

       def lowercase (		self: List[str]		)       -> Dict:




              """simple docstring"""


              with self.assertRaises((TypeError, pa.lib.ArrowInvalid)		):
                     UpperCamelCase_          =  pa.array(TypedSequence(["foo", "bar"]					,   type=Value("int64"		)		)		)

       def lowercase (		self: List[str]		)       -> List[Any]:




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence([1, 2, 3]					,   try_type=Value("int32"		)		)		)
              self.assertEqual(arr.type					,   pa.intaa()		)

       def lowercase (		self: List[str]		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence(["foo", "bar"]					,   try_type=Value("int64"		)		)		)
              self.assertEqual(arr.type					,   pa.string()		)

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




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence([[[1, 2, 3]]]					,   type=ArrayaD((1, 3)					,   "int64"		)		)		)
              self.assertEqual(arr.type					,   ArrayaDExtensionType((1, 3)					,   "int64"		)		)

       def lowercase (		self: List[Any]		)       -> Tuple:




              """simple docstring"""


              with self.assertRaises((TypeError, pa.lib.ArrowInvalid)		):
                     UpperCamelCase_          =  pa.array(TypedSequence(["foo", "bar"]					,   type=ArrayaD((1, 3)					,   "int64"		)		)		)

       def lowercase (		self: Dict		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence([[[1, 2, 3]]]					,   try_type=ArrayaD((1, 3)					,   "int64"		)		)		)
              self.assertEqual(arr.type					,   ArrayaDExtensionType((1, 3)					,   "int64"		)		)

       def lowercase (		self: List[str]		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  pa.array(TypedSequence(["foo", "bar"]					,   try_type=ArrayaD((1, 3)					,   "int64"		)		)		)
              self.assertEqual(arr.type					,   pa.string()		)



       @require_pil
       def lowercase (		self: Tuple		)       -> List[str]:




              """simple docstring"""


              import PIL.Image

              UpperCamelCase_          =  PIL.Image.fromarray(np.arange(10					,   dtype=np.uinta		).reshape(2					,   5		)		)
              with patch(
                  "datasets.arrow_writer.cast_to_python_objects"					,   side_effect=_SCREAMING_SNAKE_CASE		) as mock_cast_to_python_objects:
                     UpperCamelCase_          =  pa.array(TypedSequence([{"path": None, "bytes": B"image_bytes"}, pil_image]					,   type=Image()		)		)
                     UpperCamelCase_       , UpperCamelCase_          =  mock_cast_to_python_objects.call_args_list[-1]
                     self.assertIn("optimize_list_casting"					,   _SCREAMING_SNAKE_CASE		)
                     self.assertFalse(kwargs["optimize_list_casting"]		)




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  pa.BufferReader(UpperCamelCase_  ) if isinstance(UpperCamelCase_					,		pa.Buffer  ) else pa.memory_map(UpperCamelCase_  )
       UpperCamelCase_          =  pa.ipc.open_stream(UpperCamelCase_  )
       UpperCamelCase_          =  f.read_all()
       assert len(pa_table.to_batches()  ) == expected_num_chunks
       assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
       del pa_table




@pytest.mark.parametrize("writer_batch_size"					,		[None, 1, 10]  )
@pytest.mark.parametrize(
    "fields"					,		[None, {"col_1": pa.string(), "col_2": pa.intaa()}, {"col_1": pa.string(), "col_2": pa.intaa()}]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		int:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  pa.schema(UpperCamelCase_  ) if fields else None
       with ArrowWriter(stream=UpperCamelCase_					,		schema=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_  ) as writer:
              writer.write({"col_1": "foo", "col_2": 1}  )
              writer.write({"col_1": "bar", "col_2": 2}  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       if not fields:
              UpperCamelCase_          =  {"col_1": pa.string(), "col_2": pa.intaa()}
       assert writer._schema == pa.schema(UpperCamelCase_					,		metadata=writer._schema.metadata  )
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




def  lowerCAmelCase_ (							)     ->		str:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  Features({"labels": ClassLabel(names=["neg", "pos"]  )}  )
       with ArrowWriter(stream=UpperCamelCase_					,		features=UpperCamelCase_  ) as writer:
              writer.write({"labels": 0}  )
              writer.write({"labels": 1}  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       assert writer._schema == features.arrow_schema
       assert writer._schema.metadata == features.arrow_schema.metadata
       UpperCamelCase_          =  pa.BufferReader(output.getvalue()  )
       UpperCamelCase_          =  pa.ipc.open_stream(UpperCamelCase_  )
       UpperCamelCase_          =  f.read_all()
       UpperCamelCase_          =  pa_table.schema
       assert pa_table.num_rows == 2
       assert schema == features.arrow_schema
       assert schema.metadata == features.arrow_schema.metadata
       assert features == Features.from_arrow_schema(UpperCamelCase_  )




@pytest.mark.parametrize("writer_batch_size"					,		[None, 1, 10]  )
def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ArrowWriter(
           stream=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_					,		hash_salt="split_name"					,		check_duplicates=UpperCamelCase_					,		) as writer:
              with pytest.raises(UpperCamelCase_  ):
                     writer.write({"col_1": "foo", "col_2": 1}					,		key=[1, 2]  )
                     UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()




@pytest.mark.parametrize("writer_batch_size"					,		[None, 2, 10]  )
def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		str:
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ArrowWriter(
           stream=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_					,		hash_salt="split_name"					,		check_duplicates=UpperCamelCase_					,		) as writer:
              with pytest.raises(UpperCamelCase_  ):
                     writer.write({"col_1": "foo", "col_2": 1}					,		key=10  )
                     writer.write({"col_1": "bar", "col_2": 2}					,		key=10  )
                     UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()




@pytest.mark.parametrize("writer_batch_size"					,		[None, 2, 10]  )
def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ArrowWriter(
           stream=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_					,		hash_salt="split_name"					,		check_duplicates=UpperCamelCase_					,		) as writer:
              writer.write({"col_1": "foo", "col_2": 1}					,		key=1  )
              writer.write({"col_1": "bar", "col_2": 2}					,		key=2  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




@pytest.mark.parametrize("writer_batch_size"					,		[None, 1, 10]  )
@pytest.mark.parametrize(
    "fields"					,		[None, {"col_1": pa.string(), "col_2": pa.intaa()}, {"col_1": pa.string(), "col_2": pa.intaa()}]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[int]:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  pa.schema(UpperCamelCase_  ) if fields else None
       with ArrowWriter(stream=UpperCamelCase_					,		schema=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_  ) as writer:
              writer.write_batch({"col_1": ["foo", "bar"], "col_2": [1, 2]}  )
              writer.write_batch({"col_1": [], "col_2": []}  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       if not fields:
              UpperCamelCase_          =  {"col_1": pa.string(), "col_2": pa.intaa()}
       assert writer._schema == pa.schema(UpperCamelCase_					,		metadata=writer._schema.metadata  )
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




@pytest.mark.parametrize("writer_batch_size"					,		[None, 1, 10]  )
@pytest.mark.parametrize(
    "fields"					,		[None, {"col_1": pa.string(), "col_2": pa.intaa()}, {"col_1": pa.string(), "col_2": pa.intaa()}]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  pa.schema(UpperCamelCase_  ) if fields else None
       with ArrowWriter(stream=UpperCamelCase_					,		schema=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_  ) as writer:
              writer.write_table(pa.Table.from_pydict({"col_1": ["foo", "bar"], "col_2": [1, 2]}  )  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       if not fields:
              UpperCamelCase_          =  {"col_1": pa.string(), "col_2": pa.intaa()}
       assert writer._schema == pa.schema(UpperCamelCase_					,		metadata=writer._schema.metadata  )
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




@pytest.mark.parametrize("writer_batch_size"					,		[None, 1, 10]  )
@pytest.mark.parametrize(
    "fields"					,		[None, {"col_1": pa.string(), "col_2": pa.intaa()}, {"col_1": pa.string(), "col_2": pa.intaa()}]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[int]:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  pa.schema(UpperCamelCase_  ) if fields else None
       with ArrowWriter(stream=UpperCamelCase_					,		schema=UpperCamelCase_					,		writer_batch_size=UpperCamelCase_  ) as writer:
              writer.write_row(pa.Table.from_pydict({"col_1": ["foo"], "col_2": [1]}  )  )
              writer.write_row(pa.Table.from_pydict({"col_1": ["bar"], "col_2": [2]}  )  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       if not fields:
              UpperCamelCase_          =  {"col_1": pa.string(), "col_2": pa.intaa()}
       assert writer._schema == pa.schema(UpperCamelCase_					,		metadata=writer._schema.metadata  )
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




def  lowerCAmelCase_ (							)     ->		Any:
       with tempfile.TemporaryDirectory() as tmp_dir:
              UpperCamelCase_          =  {"col_1": pa.string(), "col_2": pa.intaa()}
              UpperCamelCase_          =  os.path.join(UpperCamelCase_					,		"test.arrow"  )
              with ArrowWriter(path=UpperCamelCase_					,		schema=pa.schema(UpperCamelCase_  )  ) as writer:
                     writer.write_batch({"col_1": ["foo", "bar"], "col_2": [1, 2]}  )
                     UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
              assert num_examples == 2
              assert num_bytes > 0
              assert writer._schema == pa.schema(UpperCamelCase_					,		metadata=writer._schema.metadata  )
              _check_output(UpperCamelCase_					,		1  )




def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Dict:
       if pa.types.is_list(UpperCamelCase_  ):
              return get_base_dtype(arr_type.value_type  )
       else:
              return arr_type




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		List[str]:
       if isinstance(lst[0]					,		UpperCamelCase_  ):
              change_first_primitive_element_in_list(lst[0]					,		UpperCamelCase_  )
       else:
              UpperCamelCase_          =  value




@pytest.mark.parametrize("optimized_int_type, expected_dtype"					,		[(None, pa.intaa()), (Value("int32"  ), pa.intaa())]  )
@pytest.mark.parametrize("sequence"					,		[[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		Dict:
       UpperCamelCase_          =  pa.array(TypedSequence(UpperCamelCase_					,		optimized_int_type=UpperCamelCase_  )  )
       assert get_base_dtype(arr.type  ) == expected_dtype




@pytest.mark.parametrize(
    "col, expected_dtype"					,		[
        ("attention_mask", pa.inta()),
        ("special_tokens_mask", pa.inta()),
        ("token_type_ids", pa.inta()),
        ("input_ids", pa.intaa()),
        ("other", pa.intaa()),
    ]					,		)
@pytest.mark.parametrize("sequence"					,		[[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		Union[str, Any]:
       # in range
       UpperCamelCase_          =  pa.array(OptimizedTypedSequence(UpperCamelCase_					,		col=UpperCamelCase_  )  )
       assert get_base_dtype(arr.type  ) == expected_dtype

       # not in range
       if col != "other":
              # avoids errors due to in-place modifications
              UpperCamelCase_          =  copy.deepcopy(UpperCamelCase_  )
              UpperCamelCase_          =  np.iinfo(expected_dtype.to_pandas_dtype()  ).max + 1
              change_first_primitive_element_in_list(UpperCamelCase_					,		UpperCamelCase_  )
              UpperCamelCase_          =  pa.array(OptimizedTypedSequence(UpperCamelCase_					,		col=UpperCamelCase_  )  )
              assert get_base_dtype(arr.type  ) == pa.intaa()




@pytest.mark.parametrize("raise_exception"					,		[False, True]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[Any]:
       UpperCamelCase_          =  str(tmp_path / "dataset-train.arrow"  )
       try:
              with ArrowWriter(path=UpperCamelCase_  ) as writer:
                     if raise_exception:
                            raise pa.lib.ArrowInvalid()
                     else:
                            writer.stream.close()
       except pa.lib.ArrowInvalid:
              pass
       finally:
              assert writer.stream.closed




def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Optional[int]:
       UpperCamelCase_          =  "mock://dataset-train.arrow"
       with ArrowWriter(path=UpperCamelCase_					,		storage_options=mockfs.storage_options  ) as writer:
              assert isinstance(writer._fs					,		type(UpperCamelCase_  )  )
              assert writer._fs.storage_options == mockfs.storage_options
              writer.write({"col_1": "foo", "col_2": 1}  )
              writer.write({"col_1": "bar", "col_2": 2}  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       assert mockfs.exists(UpperCamelCase_  )




def  lowerCAmelCase_ (							)     ->		Optional[int]:
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ParquetWriter(stream=UpperCamelCase_  ) as writer:
              writer.write({"col_1": "foo", "col_2": 1}  )
              writer.write({"col_1": "bar", "col_2": 2}  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       UpperCamelCase_          =  pa.BufferReader(output.getvalue()  )
       UpperCamelCase_          =  pq.read_table(UpperCamelCase_  )
       assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}




@require_pil
@pytest.mark.parametrize("embed_local_files"					,		[False, True]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[Any]:
       import PIL.Image

       UpperCamelCase_          =  str(tmp_path / "test_image_rgb.jpg"  )
       PIL.Image.fromarray(np.zeros((5, 5)					,		dtype=np.uinta  )  ).save(UpperCamelCase_					,		format="png"  )
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ParquetWriter(
           stream=UpperCamelCase_					,		features=Features({"image": Image()}  )					,		embed_local_files=UpperCamelCase_  ) as writer:
              writer.write({"image": image_path}  )
              writer.finalize()
       UpperCamelCase_          =  pa.BufferReader(output.getvalue()  )
       UpperCamelCase_          =  pq.read_table(UpperCamelCase_  )
       UpperCamelCase_          =  pa_table.to_pydict()
       if embed_local_files:
              assert isinstance(out["image"][0]["path"]					,		UpperCamelCase_  )
              with open(UpperCamelCase_					,		"rb"  ) as f:
                     assert out["image"][0]["bytes"] == f.read()
       else:
              assert out["image"][0]["path"] == image_path
              assert out["image"][0]["bytes"] is None







def  lowerCAmelCase_ (							)     ->		int:
       UpperCamelCase_          =  pa.schema([pa.field("col_1"					,		pa.string()					,		nullable=UpperCamelCase_  )]  )
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ArrowWriter(stream=UpperCamelCase_  ) as writer:
              writer._build_writer(inferred_schema=UpperCamelCase_  )
              assert writer._schema == pa.schema([pa.field("col_1"					,		pa.string()  )]  )








import unittest

from transformers import AutoTokenizer, FalconConfig, 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 (
         FalconForCausalLM,
         FalconForQuestionAnswering,
         FalconForSequenceClassification,
         FalconForTokenClassification,
         FalconModel,
     )



class 							_UpperCamelCase					:

       def __init__(		self: Dict					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]=3					,   _SCREAMING_SNAKE_CASE: Any=7					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=True					,   _SCREAMING_SNAKE_CASE: Optional[int]=True					,   _SCREAMING_SNAKE_CASE: List[Any]=False					,   _SCREAMING_SNAKE_CASE: str=True					,   _SCREAMING_SNAKE_CASE: Any=99					,   _SCREAMING_SNAKE_CASE: List[str]=32					,   _SCREAMING_SNAKE_CASE: List[str]=5					,   _SCREAMING_SNAKE_CASE: Optional[int]=4					,   _SCREAMING_SNAKE_CASE: str=37					,   _SCREAMING_SNAKE_CASE: Optional[int]="gelu"					,   _SCREAMING_SNAKE_CASE: str=0.1					,   _SCREAMING_SNAKE_CASE: Any=0.1					,   _SCREAMING_SNAKE_CASE: Optional[Any]=512					,   _SCREAMING_SNAKE_CASE: Optional[Any]=16					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=2					,   _SCREAMING_SNAKE_CASE: Tuple=0.02					,   _SCREAMING_SNAKE_CASE: Optional[Any]=3					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=4					,   _SCREAMING_SNAKE_CASE: Tuple=None					,   )       -> Tuple:




              """simple docstring"""


              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

       def lowercase (		self: Any		)       -> Any:




              """simple docstring"""


              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
              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, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

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




              """simple docstring"""


              return FalconConfig(
                  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					,   hidden_dropout_prob=self.hidden_dropout_prob					,   attention_probs_dropout_prob=self.attention_probs_dropout_prob					,   max_position_embeddings=self.max_position_embeddings					,   type_vocab_size=self.type_vocab_size					,   is_decoder=_SCREAMING_SNAKE_CASE					,   initializer_range=self.initializer_range					,   pad_token_id=1					,   new_decoder_architecture=_SCREAMING_SNAKE_CASE					,   )

       def lowercase (		self: str					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: Dict					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: int					,   _SCREAMING_SNAKE_CASE: Dict					,   _SCREAMING_SNAKE_CASE: Union[str, Any]		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  FalconModel(config=_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE		)
              self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size)		)

       def lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Union[str, Any]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Tuple					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: str					,   _SCREAMING_SNAKE_CASE: Tuple					,   )       -> str:




              """simple docstring"""


              UpperCamelCase_          =  True
              UpperCamelCase_          =  FalconModel(_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(
                  _SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   encoder_hidden_states=_SCREAMING_SNAKE_CASE					,   encoder_attention_mask=_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  model(
                  _SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   encoder_hidden_states=_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE		)
              self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size)		)

       def lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Union[str, Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Union[str, Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: int					,   )       -> int:




              """simple docstring"""


              UpperCamelCase_          =  FalconForCausalLM(config=_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   labels=_SCREAMING_SNAKE_CASE		)
              self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.vocab_size)		)

       def lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: int					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: Union[str, Any]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Optional[Any]					,   _SCREAMING_SNAKE_CASE: Tuple					,   _SCREAMING_SNAKE_CASE: Union[str, Any]					,   )       -> List[Any]:




              """simple docstring"""


              UpperCamelCase_          =  True
              UpperCamelCase_          =  True
              UpperCamelCase_          =  FalconForCausalLM(config=_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()

              # first forward pass
              UpperCamelCase_          =  model(
                  _SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   encoder_hidden_states=_SCREAMING_SNAKE_CASE					,   encoder_attention_mask=_SCREAMING_SNAKE_CASE					,   use_cache=_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  outputs.past_key_values

              # create hypothetical multiple next token and extent to next_input_ids
              UpperCamelCase_          =  ids_tensor((self.batch_size, 3)					,   config.vocab_size		)
              UpperCamelCase_          =  ids_tensor((self.batch_size, 3)					,   vocab_size=2		)

              # append to next input_ids and
              UpperCamelCase_          =  torch.cat([input_ids, next_tokens]					,   dim=-1		)
              UpperCamelCase_          =  torch.cat([input_mask, next_mask]					,   dim=-1		)

              UpperCamelCase_          =  model(
                  _SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   encoder_hidden_states=_SCREAMING_SNAKE_CASE					,   encoder_attention_mask=_SCREAMING_SNAKE_CASE					,   output_hidden_states=_SCREAMING_SNAKE_CASE					,   )["hidden_states"][0]
              UpperCamelCase_          =  model(
                  _SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   encoder_hidden_states=_SCREAMING_SNAKE_CASE					,   encoder_attention_mask=_SCREAMING_SNAKE_CASE					,   past_key_values=_SCREAMING_SNAKE_CASE					,   output_hidden_states=_SCREAMING_SNAKE_CASE					,   )["hidden_states"][0]

              # select random slice
              UpperCamelCase_          =  ids_tensor((1,)					,   output_from_past.shape[-1]		).item()
              UpperCamelCase_          =  output_from_no_past[:, -3:, random_slice_idx].detach()
              UpperCamelCase_          =  output_from_past[:, :, random_slice_idx].detach()

              self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

              # test that outputs are equal for slice
              self.parent.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   atol=1e-3		)		)



       def lowercase (		self: Tuple		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  self.prepare_config_and_inputs()
              (
                  (
                  UpperCamelCase_
              )       , (
                  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_  ,						lowerCAmelCase_  ,						unittest.TestCase     ):
       _UpperCamelCase				:					Optional[Any]         =   (
           (
               FalconModel,
               FalconForCausalLM,
               FalconForSequenceClassification,
               FalconForTokenClassification,
               FalconForQuestionAnswering,
           )
           if is_torch_available()
           else ()
       )
       _UpperCamelCase				:					Union[str, Any]         =   (FalconForCausalLM,) if is_torch_available() else ()
       _UpperCamelCase				:					Optional[int]         =   (
           {
               '''feature-extraction''': FalconModel,
               '''text-classification''': FalconForSequenceClassification,
               '''text-generation''': FalconForCausalLM,
               '''question-answering''': FalconForQuestionAnswering,
               '''token-classification''': FalconForTokenClassification,
               '''zero-shot''': FalconForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       _UpperCamelCase				:					Any         =   False
       _UpperCamelCase				:					int         =   False

       def lowercase (		self: str		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  FalconModelTester(self		)
              UpperCamelCase_          =  ConfigTester(self					,   config_class=_SCREAMING_SNAKE_CASE					,   hidden_size=37		)

       def lowercase (		self: Optional[Any]		)       -> Optional[int]:




              """simple docstring"""


              self.config_tester.run_common_tests()

       def lowercase (		self: str		)       -> Any:




              """simple docstring"""


              UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: List[Any]		)       -> Dict:




              """simple docstring"""


              UpperCamelCase_       , *UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs()
              for alibi in [True, False]:
                     UpperCamelCase_          =  alibi
                     self.model_tester.create_and_check_model(_SCREAMING_SNAKE_CASE					,   *_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Any		)       -> str:




              """simple docstring"""


              UpperCamelCase_       , UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs_for_common()
              UpperCamelCase_          =  3
              UpperCamelCase_          =  input_dict["input_ids"]
              UpperCamelCase_          =  input_ids.ne(1		).to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  ids_tensor([self.model_tester.batch_size]					,   self.model_tester.type_sequence_label_size		)
              UpperCamelCase_          =  FalconForSequenceClassification(_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   labels=_SCREAMING_SNAKE_CASE		)
              self.assertEqual(result.logits.shape					,   (self.model_tester.batch_size, self.model_tester.num_labels)		)

       def lowercase (		self: Dict		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_       , UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs_for_common()
              UpperCamelCase_          =  3
              UpperCamelCase_          =  "single_label_classification"
              UpperCamelCase_          =  input_dict["input_ids"]
              UpperCamelCase_          =  input_ids.ne(1		).to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  ids_tensor([self.model_tester.batch_size]					,   self.model_tester.type_sequence_label_size		)
              UpperCamelCase_          =  FalconForSequenceClassification(_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   labels=_SCREAMING_SNAKE_CASE		)
              self.assertEqual(result.logits.shape					,   (self.model_tester.batch_size, self.model_tester.num_labels)		)

       def lowercase (		self: List[str]		)       -> int:




              """simple docstring"""


              UpperCamelCase_       , UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs_for_common()
              UpperCamelCase_          =  input_dict["input_ids"]
              UpperCamelCase_          =  FalconForCausalLM(_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   use_cache=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  input_ids.shape[0]
              UpperCamelCase_          =  model._convert_to_rw_cache(result.past_key_values		)
              UpperCamelCase_          =  model._convert_cache_to_standard_format(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
              for layer in range(len(_SCREAMING_SNAKE_CASE		)		):
                     for tensor_idx in range(2		):
                            self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3		)
                            self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4		)
                            self.assertTrue(
                                torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]		)		)

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




              """simple docstring"""


              UpperCamelCase_       , UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs_for_common()
              UpperCamelCase_          =  3
              UpperCamelCase_          =  "multi_label_classification"
              UpperCamelCase_          =  input_dict["input_ids"]
              UpperCamelCase_          =  input_ids.ne(1		).to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  ids_tensor(
                  [self.model_tester.batch_size, config.num_labels]					,   self.model_tester.type_sequence_label_size		).to(torch.float		)
              UpperCamelCase_          =  FalconForSequenceClassification(_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()
              UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   attention_mask=_SCREAMING_SNAKE_CASE					,   labels=_SCREAMING_SNAKE_CASE		)
              self.assertEqual(result.logits.shape					,   (self.model_tester.batch_size, self.model_tester.num_labels)		)



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




              """simple docstring"""


              for model_class in self.all_generative_model_classes:
                     UpperCamelCase_       , UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs_for_common()

                     # If it doesn't support cache, pass the test
                     if not hasattr(_SCREAMING_SNAKE_CASE					,   "use_cache"		):
                            return

                     UpperCamelCase_          =  model_class(_SCREAMING_SNAKE_CASE		).to(_SCREAMING_SNAKE_CASE		)
                     if "use_cache" not in inputs:
                            UpperCamelCase_          =  True
                     UpperCamelCase_          =  model(**_SCREAMING_SNAKE_CASE		)

                     # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
                     if "past_key_values" not in outputs:
                            return

                     UpperCamelCase_          =  (
                         getattr(_SCREAMING_SNAKE_CASE					,   "decoder_layers"					,   _SCREAMING_SNAKE_CASE		)
                         or getattr(_SCREAMING_SNAKE_CASE					,   "num_decoder_layers"					,   _SCREAMING_SNAKE_CASE		)
                         or config.num_hidden_layers
                     )
                     UpperCamelCase_          =  getattr(_SCREAMING_SNAKE_CASE					,   "num_kv_heads"					,   config.num_attention_heads		)
                     UpperCamelCase_          =  getattr(_SCREAMING_SNAKE_CASE					,   "d_model"					,   config.hidden_size		)
                     UpperCamelCase_          =  embed_dim // num_attention_heads

                     UpperCamelCase_          =  outputs["past_key_values"]
                     self.assertEqual(len(_SCREAMING_SNAKE_CASE		)					,   _SCREAMING_SNAKE_CASE		)

                     UpperCamelCase_       , UpperCamelCase_          =  inputs["input_ids"].shape
                     for i in range(_SCREAMING_SNAKE_CASE		):
                            if config.new_decoder_architecture:
                                   UpperCamelCase_          =  config.num_attention_heads
                            elif config.multi_query:
                                   UpperCamelCase_          =  1
                            self.assertEqual(len(past_kv[0]		)					,   2		)  # K V for the decoder = 2
                            self.assertEqual(
                                past_kv[i][0].shape					,   (batch_size, num_attention_heads, seq_length, per_head_embed_dim)		)
                            self.assertEqual(
                                past_kv[i][1].shape					,   (batch_size, num_attention_heads, seq_length, per_head_embed_dim)		)



@require_torch
class 							_UpperCamelCase					(	unittest.TestCase     ):

       @slow
       def lowercase (		self: Tuple		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"		)
              UpperCamelCase_          =  FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"		)
              model.eval()
              model.to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  tokenizer("My favorite food is"					,   return_tensors="pt"		).to(_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  (
                  "My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."
              )

              UpperCamelCase_          =  model.generate(**_SCREAMING_SNAKE_CASE					,   do_sample=_SCREAMING_SNAKE_CASE					,   max_new_tokens=19		)
              UpperCamelCase_          =  tokenizer.batch_decode(_SCREAMING_SNAKE_CASE		)[0]

              self.assertEqual(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

       @slow
       def lowercase (		self: Dict		)       -> Dict:




              """simple docstring"""


              for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
                     UpperCamelCase_          =  AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  FalconForCausalLM.from_pretrained(_SCREAMING_SNAKE_CASE		)
                     model.eval()
                     model.to(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  tokenizer("My favorite food is"					,   return_tensors="pt"		).to(_SCREAMING_SNAKE_CASE		)

                     # We just test that these run without errors - the models are randomly initialized
                     # and so the actual text outputs will be garbage
                     model.generate(**_SCREAMING_SNAKE_CASE					,   do_sample=_SCREAMING_SNAKE_CASE					,   max_new_tokens=4		)
                     model.generate(**_SCREAMING_SNAKE_CASE					,   do_sample=_SCREAMING_SNAKE_CASE					,   max_new_tokens=4		)
                     model.generate(**_SCREAMING_SNAKE_CASE					,   num_beams=2					,   max_new_tokens=4		)



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




              """simple docstring"""


              with torch.no_grad():
                     for repo in [
                         "Rocketknight1/falcon-rw-1b",
                         "Rocketknight1/tiny-random-falcon-7b",
                         "Rocketknight1/tiny-random-falcon-40b",
                     ]:
                            UpperCamelCase_          =  AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE		)
                            UpperCamelCase_          =  FalconForCausalLM.from_pretrained(_SCREAMING_SNAKE_CASE		)
                            model.eval()
                            model.to(device=_SCREAMING_SNAKE_CASE		)
                            UpperCamelCase_          =  tokenizer("My favorite food is"					,   return_tensors="pt"		).to(_SCREAMING_SNAKE_CASE		)

                            # Test results are the same with and without cache
                            UpperCamelCase_          =  model.generate(**_SCREAMING_SNAKE_CASE					,   do_sample=_SCREAMING_SNAKE_CASE					,   max_new_tokens=20					,   use_cache=_SCREAMING_SNAKE_CASE		)
                            UpperCamelCase_          =  model.generate(**_SCREAMING_SNAKE_CASE					,   do_sample=_SCREAMING_SNAKE_CASE					,   max_new_tokens=20					,   use_cache=_SCREAMING_SNAKE_CASE		)
                            self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0		)

