'''simple docstring'''






import contextlib
import csv
import json
import os
import sqlitea
import tarfile
import textwrap
import zipfile

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

import datasets
import datasets.config






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					)							->	Optional[int]:
  __UpperCAmelCase:					int		= 1_0
  __UpperCAmelCase:					List[Any]		= datasets.Features(
      {
          """tokens""": datasets.Sequence(datasets.Value("""string"""     )     ),
          """labels""": datasets.Sequence(datasets.ClassLabel(names=["""negative""", """positive"""]     )     ),
          """answers""": datasets.Sequence(
              {
                  """text""": datasets.Value("""string"""     ),
                  """answer_start""": datasets.Value("""int32"""     ),
              }     ),
          """id""": datasets.Value("""int64"""     ),
      }     )
  __UpperCAmelCase:					str		= datasets.Dataset.from_dict(
      {
          """tokens""": [["""foo"""] * 5] * n,
          """labels""": [[1] * 5] * n,
          """answers""": [{"""answer_start""": [9_7], """text""": ["""1976"""]}] * 1_0,
          """id""": list(range(_lowercase     )     ),
      }   ,    features=_lowercase   ,    )
  return dataset






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Optional[Any]   ,    _lowercase    :				Optional[int]     )							->	Union[str, Any]:
  __UpperCAmelCase:					Union[str, Any]		= str(tmp_path_factory.mktemp("""data"""     ) / """file.arrow"""     )
  dataset.map(cache_file_name=_lowercase     )
  return filename


# FILE_CONTENT + files


SCREAMING_SNAKE_CASE_	   =			'\\n    Text data.\n    Second line of data.'






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[Any]     )							->	Tuple:
  __UpperCAmelCase:					Tuple		= tmp_path_factory.mktemp("""data"""     ) / """file.txt"""
  __UpperCAmelCase:					str		= FILE_CONTENT
  with open(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase     )
  return filename






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Union[str, Any]     )							->	Tuple:
  import bza

  __UpperCAmelCase:					Tuple		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.bz2"""
  __UpperCAmelCase:					str		= bytes(_lowercase   ,    """utf-8"""     )
  with bza.open(_lowercase   ,    """wb"""     ) as f:
    f.write(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Any     )							->	List[Any]:
  import gzip

  __UpperCAmelCase:					List[Any]		= str(tmp_path_factory.mktemp("""data"""     ) / """file.txt.gz"""     )
  __UpperCAmelCase:					Optional[int]		= bytes(_lowercase   ,    """utf-8"""     )
  with gzip.open(_lowercase   ,    """wb"""     ) as f:
    f.write(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int     )							->	Union[str, Any]:
  if datasets.config.LZ4_AVAILABLE:
    import lza.frame

    __UpperCAmelCase:					List[str]		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.lz4"""
    __UpperCAmelCase:					Union[str, Any]		= bytes(_lowercase   ,    """utf-8"""     )
    with lza.frame.open(_lowercase   ,    """wb"""     ) as f:
      f.write(_lowercase     )
    return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Optional[int]   ,    _lowercase    :				Any     )							->	Any:
  if datasets.config.PY7ZR_AVAILABLE:
    import pyazr

    __UpperCAmelCase:					List[str]		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.7z"""
    with pyazr.SevenZipFile(_lowercase   ,    """w"""     ) as archive:
      archive.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
    return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Any   ,    _lowercase    :				Tuple     )							->	int:
  import tarfile

  __UpperCAmelCase:					str		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.tar"""
  with tarfile.TarFile(_lowercase   ,    """w"""     ) as f:
    f.add(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Optional[Any]     )							->	Optional[int]:
  import lzma

  __UpperCAmelCase:					Union[str, Any]		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.xz"""
  __UpperCAmelCase:					Optional[int]		= bytes(_lowercase   ,    """utf-8"""     )
  with lzma.open(_lowercase   ,    """wb"""     ) as f:
    f.write(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Dict   ,    _lowercase    :				Dict     )							->	List[str]:
  import zipfile

  __UpperCAmelCase:					int		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Optional[int]     )							->	str:
  if datasets.config.ZSTANDARD_AVAILABLE:
    import zstandard as zstd

    __UpperCAmelCase:					List[str]		= tmp_path_factory.mktemp("""data"""     ) / """file.txt.zst"""
    __UpperCAmelCase:					List[Any]		= bytes(_lowercase   ,    """utf-8"""     )
    with zstd.open(_lowercase   ,    """wb"""     ) as f:
      f.write(_lowercase     )
    return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Dict     )							->	Union[str, Any]:
  __UpperCAmelCase:					Union[str, Any]		= tmp_path_factory.mktemp("""data"""     ) / """file.xml"""
  __UpperCAmelCase:					int		= textwrap.dedent(
      """\
    <?xml version=\"1.0\" encoding=\"UTF-8\" ?>
    <tmx version=\"1.4\">
      <header segtype=\"sentence\" srclang=\"ca\" />
      <body>
        <tu>
          <tuv xml:lang=\"ca\"><seg>Contingut 1</seg></tuv>
          <tuv xml:lang=\"en\"><seg>Content 1</seg></tuv>
        </tu>
        <tu>
          <tuv xml:lang=\"ca\"><seg>Contingut 2</seg></tuv>
          <tuv xml:lang=\"en\"><seg>Content 2</seg></tuv>
        </tu>
        <tu>
          <tuv xml:lang=\"ca\"><seg>Contingut 3</seg></tuv>
          <tuv xml:lang=\"en\"><seg>Content 3</seg></tuv>
        </tu>
        <tu>
          <tuv xml:lang=\"ca\"><seg>Contingut 4</seg></tuv>
          <tuv xml:lang=\"en\"><seg>Content 4</seg></tuv>
        </tu>
        <tu>
          <tuv xml:lang=\"ca\"><seg>Contingut 5</seg></tuv>
          <tuv xml:lang=\"en\"><seg>Content 5</seg></tuv>
        </tu>
      </body>
    </tmx>"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase     )
  return filename


SCREAMING_SNAKE_CASE_	   =			[
    {'col_1': '0', 'col_2': 0, 'col_3': 0.0},
    {'col_1': '1', 'col_2': 1, 'col_3': 1.0},
    {'col_1': '2', 'col_2': 2, 'col_3': 2.0},
    {'col_1': '3', 'col_2': 3, 'col_3': 3.0},
]
SCREAMING_SNAKE_CASE_	   =			[
    {'col_1': '4', 'col_2': 4, 'col_3': 4.0},
    {'col_1': '5', 'col_2': 5, 'col_3': 5.0},
]
SCREAMING_SNAKE_CASE_	   =			{
    'col_1': ['0', '1', '2', '3'],
    'col_2': [0, 1, 2, 3],
    'col_3': [0.0, 1.0, 2.0, 3.0],
}

SCREAMING_SNAKE_CASE_	   =			[
    {'col_3': 0.0, 'col_1': '0', 'col_2': 0},
    {'col_3': 1.0, 'col_1': '1', 'col_2': 1},
]

SCREAMING_SNAKE_CASE_	   =			[
    {'col_1': 's0', 'col_2': 0, 'col_3': 0.0},
    {'col_1': 's1', 'col_2': 1, 'col_3': 1.0},
    {'col_1': 's2', 'col_2': 2, 'col_3': 2.0},
    {'col_1': 's3', 'col_2': 3, 'col_3': 3.0},
]






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					)							->	Any:
  return DATA_DICT_OF_LISTS






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int     )							->	int:
  __UpperCAmelCase:					str		= datasets.Dataset.from_dict(_lowercase     )
  __UpperCAmelCase:					List[str]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.arrow"""     )
  dataset.map(cache_file_name=_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Dict     )							->	Any:
  __UpperCAmelCase:					Optional[int]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.sqlite"""     )
  with contextlib.closing(sqlitea.connect(_lowercase     )     ) as con:
    __UpperCAmelCase:					str		= con.cursor()
    cur.execute("""CREATE TABLE dataset(col_1 text, col_2 int, col_3 real)"""     )
    for item in DATA:
      cur.execute("""INSERT INTO dataset(col_1, col_2, col_3) VALUES (?, ?, ?)"""   ,    tuple(item.values()     )     )
    con.commit()
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Union[str, Any]     )							->	Optional[Any]:
  __UpperCAmelCase:					int		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.csv"""     )
  with open(_lowercase   ,    """w"""   ,    newline=""""""     ) as f:
    __UpperCAmelCase:					Optional[Any]		= csv.DictWriter(_lowercase   ,    fieldnames=["""col_1""", """col_2""", """col_3"""]     )
    writer.writeheader()
    for item in DATA:
      writer.writerow(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[str]     )							->	List[Any]:
  __UpperCAmelCase:					List[Any]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset2.csv"""     )
  with open(_lowercase   ,    """w"""   ,    newline=""""""     ) as f:
    __UpperCAmelCase:					str		= csv.DictWriter(_lowercase   ,    fieldnames=["""col_1""", """col_2""", """col_3"""]     )
    writer.writeheader()
    for item in DATA:
      writer.writerow(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Tuple   ,    _lowercase    :				Optional[Any]     )							->	str:
  import bza

  __UpperCAmelCase:					Optional[int]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.csv.bz2"""
  with open(_lowercase   ,    """rb"""     ) as f:
    __UpperCAmelCase:					Union[str, Any]		= f.read()
  # data = bytes(FILE_CONTENT, "utf-8")
  with bza.open(_lowercase   ,    """wb"""     ) as f:
    f.write(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Union[str, Any]   ,    _lowercase    :				str   ,    _lowercase    :				List[str]     )							->	Tuple:
  __UpperCAmelCase:					int		= tmp_path_factory.mktemp("""data"""     ) / """dataset.csv.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[Any]   ,    _lowercase    :				Any   ,    _lowercase    :				Optional[int]     )							->	List[Any]:
  __UpperCAmelCase:					Optional[Any]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.csv.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename(csv_path.replace(""".csv"""   ,    """.CSV"""     )     )     )
    f.write(_lowercase   ,    arcname=os.path.basename(csva_path.replace(""".csv"""   ,    """.CSV"""     )     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				str   ,    _lowercase    :				List[Any]   ,    _lowercase    :				List[Any]     )							->	Any:
  __UpperCAmelCase:					Dict		= tmp_path_factory.mktemp("""data"""     ) / """dataset_with_dir.csv.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.join("""main_dir"""   ,    os.path.basename(_lowercase     )     )     )
    f.write(_lowercase   ,    arcname=os.path.join("""main_dir"""   ,    os.path.basename(_lowercase     )     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Optional[int]     )							->	List[str]:
  __UpperCAmelCase:					int		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.parquet"""     )
  __UpperCAmelCase:					List[Any]		= pa.schema(
      {
          """col_1""": pa.string(),
          """col_2""": pa.intaa(),
          """col_3""": pa.floataa(),
      }     )
  with open(_lowercase   ,    """wb"""     ) as f:
    __UpperCAmelCase:					int		= pq.ParquetWriter(_lowercase   ,    schema=_lowercase     )
    __UpperCAmelCase:					List[str]		= pa.Table.from_pydict({k: [DATA[i][k] for i in range(len(_lowercase     )     )] for k in DATA[0]}   ,    schema=_lowercase     )
    writer.write_table(_lowercase     )
    writer.close()
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Dict     )							->	List[str]:
  __UpperCAmelCase:					Tuple		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.json"""     )
  __UpperCAmelCase:					List[str]		= {"""data""": DATA}
  with open(_lowercase   ,    """w"""     ) as f:
    json.dump(_lowercase   ,    _lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				str     )							->	Any:
  __UpperCAmelCase:					int		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.json"""     )
  __UpperCAmelCase:					Any		= {"""data""": DATA_DICT_OF_LISTS}
  with open(_lowercase   ,    """w"""     ) as f:
    json.dump(_lowercase   ,    _lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Any     )							->	List[Any]:
  __UpperCAmelCase:					Optional[int]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.jsonl"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    for item in DATA:
      f.write(json.dumps(_lowercase     ) + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Tuple     )							->	Any:
  __UpperCAmelCase:					Optional[Any]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset2.jsonl"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    for item in DATA:
      f.write(json.dumps(_lowercase     ) + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Tuple     )							->	Dict:
  __UpperCAmelCase:					int		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset_312.jsonl"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    for item in DATA_312:
      f.write(json.dumps(_lowercase     ) + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int     )							->	Any:
  __UpperCAmelCase:					Optional[int]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset-str.jsonl"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    for item in DATA_STR:
      f.write(json.dumps(_lowercase     ) + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int   ,    _lowercase    :				Dict     )							->	Dict:
  import gzip

  __UpperCAmelCase:					str		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.txt.gz"""     )
  with open(_lowercase   ,    """rb"""     ) as orig_file:
    with gzip.open(_lowercase   ,    """wb"""     ) as zipped_file:
      zipped_file.writelines(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int   ,    _lowercase    :				List[Any]     )							->	Dict:
  import gzip

  __UpperCAmelCase:					int		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.jsonl.gz"""     )
  with open(_lowercase   ,    """rb"""     ) as orig_file:
    with gzip.open(_lowercase   ,    """wb"""     ) as zipped_file:
      zipped_file.writelines(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[str]   ,    _lowercase    :				str   ,    _lowercase    :				List[str]     )							->	Any:
  __UpperCAmelCase:					str		= tmp_path_factory.mktemp("""data"""     ) / """dataset.jsonl.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Union[str, Any]   ,    _lowercase    :				Optional[Any]   ,    _lowercase    :				str   ,    _lowercase    :				Tuple     )							->	Tuple:
  __UpperCAmelCase:					Any		= tmp_path_factory.mktemp("""data"""     ) / """dataset_nested.jsonl.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.join("""nested"""   ,    os.path.basename(_lowercase     )     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[str]   ,    _lowercase    :				Optional[int]   ,    _lowercase    :				List[str]     )							->	Dict:
  __UpperCAmelCase:					int		= tmp_path_factory.mktemp("""data"""     ) / """dataset_with_dir.jsonl.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.join("""main_dir"""   ,    os.path.basename(_lowercase     )     )     )
    f.write(_lowercase   ,    arcname=os.path.join("""main_dir"""   ,    os.path.basename(_lowercase     )     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int   ,    _lowercase    :				int   ,    _lowercase    :				List[str]     )							->	Dict:
  __UpperCAmelCase:					Optional[int]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.jsonl.tar"""
  with tarfile.TarFile(_lowercase   ,    """w"""     ) as f:
    f.add(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
    f.add(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Tuple   ,    _lowercase    :				Optional[Any]   ,    _lowercase    :				Optional[Any]   ,    _lowercase    :				Optional[int]     )							->	Union[str, Any]:
  __UpperCAmelCase:					Optional[int]		= tmp_path_factory.mktemp("""data"""     ) / """dataset_nested.jsonl.tar"""
  with tarfile.TarFile(_lowercase   ,    """w"""     ) as f:
    f.add(_lowercase   ,    arcname=os.path.join("""nested"""   ,    os.path.basename(_lowercase     )     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int     )							->	Dict:
  __UpperCAmelCase:					Tuple		= ["""0""", """1""", """2""", """3"""]
  __UpperCAmelCase:					Optional[int]		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset.txt"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    for item in data:
      f.write(item + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Tuple     )							->	List[str]:
  __UpperCAmelCase:					List[str]		= ["""0""", """1""", """2""", """3"""]
  __UpperCAmelCase:					str		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset2.txt"""     )
  with open(_lowercase   ,    """w"""     ) as f:
    for item in data:
      f.write(item + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[str]     )							->	List[str]:
  __UpperCAmelCase:					Tuple		= ["""0""", """1""", """2""", """3"""]
  __UpperCAmelCase:					Optional[int]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.abc"""
  with open(_lowercase   ,    """w"""     ) as f:
    for item in data:
      f.write(item + """\n"""     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[Any]   ,    _lowercase    :				int   ,    _lowercase    :				Any     )							->	int:
  __UpperCAmelCase:					Optional[int]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.text.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				int   ,    _lowercase    :				Dict   ,    _lowercase    :				Optional[int]     )							->	int:
  __UpperCAmelCase:					Dict		= tmp_path_factory.mktemp("""data"""     ) / """dataset_with_dir.text.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.join("""main_dir"""   ,    os.path.basename(_lowercase     )     )     )
    f.write(_lowercase   ,    arcname=os.path.join("""main_dir"""   ,    os.path.basename(_lowercase     )     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[Any]   ,    _lowercase    :				Tuple   ,    _lowercase    :				int     )							->	List[Any]:
  __UpperCAmelCase:					Optional[Any]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.ext.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename("""unsupported.ext"""     )     )
    f.write(_lowercase   ,    arcname=os.path.basename("""unsupported_2.ext"""     )     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[str]     )							->	Dict:
  __UpperCAmelCase:					List[str]		= """\n""".join(["""First""", """Second\u2029with Unicode new line""", """Third"""]     )
  __UpperCAmelCase:					Dict		= str(tmp_path_factory.mktemp("""data"""     ) / """dataset_with_unicode_new_lines.txt"""     )
  with open(_lowercase   ,    """w"""   ,    encoding="""utf-8"""     ) as f:
    f.write(_lowercase     )
  return path






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					)							->	str:
  return os.path.join("""tests"""   ,    """features"""   ,    """data"""   ,    """test_image_rgb.jpg"""     )






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					)							->	Union[str, Any]:
  return os.path.join("""tests"""   ,    """features"""   ,    """data"""   ,    """test_audio_44100.wav"""     )






@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				List[Any]   ,    _lowercase    :				List[Any]     )							->	Dict:
  __UpperCAmelCase:					Optional[int]		= tmp_path_factory.mktemp("""data"""     ) / """dataset.img.zip"""
  with zipfile.ZipFile(_lowercase   ,    """w"""     ) as f:
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     )     )
    f.write(_lowercase   ,    arcname=os.path.basename(_lowercase     ).replace(""".jpg"""   ,    """2.jpg"""     )     )
  return path







@pytest.fixture(scope="""session"""     )
def  UpperCamelCase__					(					_lowercase    :				Optional[Any]     )							->	str:
  __UpperCAmelCase:					Dict		= tmp_path_factory.mktemp("""data_dir"""     )

  (data_dir / "subdir").mkdir()
  with open(data_dir / """subdir""" / """train.txt"""   ,    """w"""     ) as f:
    f.write("""foo\n""" * 1_0     )
  with open(data_dir / """subdir""" / """test.txt"""   ,    """w"""     ) as f:
    f.write("""bar\n""" * 1_0     )
  # hidden file
  with open(data_dir / """subdir""" / """.test.txt"""   ,    """w"""     ) as f:
    f.write("""bar\n""" * 1_0     )

  # hidden directory
  (data_dir / ".subdir").mkdir()
  with open(data_dir / """.subdir""" / """train.txt"""   ,    """w"""     ) as f:
    f.write("""foo\n""" * 1_0     )
  with open(data_dir / """.subdir""" / """test.txt"""   ,    """w"""     ) as f:
    f.write("""bar\n""" * 1_0     )

  return data_dir
'''simple docstring'''






import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


SCREAMING_SNAKE_CASE_	   =			get_logger(__name__)


SCREAMING_SNAKE_CASE_	   =			R'\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n'




class    a :



  """simple docstring"""







  @add_start_docstrings(snake_case_					)
  def __call__(							self   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    raise NotImplementedError(
        F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''					)




class    a :



  """simple docstring"""







  @add_start_docstrings(snake_case_					)
  def __call__(							self   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    raise NotImplementedError(
        F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''					)




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  @add_start_docstrings(snake_case_					)
  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    **snake_case_					):

    '''simple docstring'''







    for processor in self:
      __UpperCAmelCase:					str		= inspect.signature(processor.__call__					).parameters
      if len(snake_case_					) > 3:
        if not all(arg in kwargs for arg in list(function_args.keys()					)[2:]					):
          raise ValueError(
              F'''Make sure that all the required parameters: {list(function_args.keys()					)} for '''
              F'''{processor.__class__} are passed to the logits processor.'''					)
        __UpperCAmelCase:					List[Any]		= processor(snake_case_   ,    snake_case_   ,    snake_case_   ,    **snake_case_					)
      else:
        __UpperCAmelCase:					Optional[Any]		= processor(snake_case_   ,    snake_case_   ,    snake_case_					)
    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_					):

    '''simple docstring'''







    if not isinstance(snake_case_   ,    snake_case_					) or not (temperature > 0):
      raise ValueError(F'''`temperature` has to be a strictly positive float, but is {temperature}'''					)

    __UpperCAmelCase:					Any		= temperature







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Tuple		= scores / self.temperature
    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_ = -float("""Inf"""					)   ,    snake_case_ = 1					):

    '''simple docstring'''







    if not isinstance(snake_case_   ,    snake_case_					) or (top_p < 0 or top_p > 1.0):
      raise ValueError(F'''`top_p` has to be a float > 0 and < 1, but is {top_p}'''					)
    if not isinstance(snake_case_   ,    snake_case_					) or (min_tokens_to_keep < 1):
      raise ValueError(F'''`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}'''					)

    __UpperCAmelCase:					int		= top_p
    __UpperCAmelCase:					str		= filter_value
    __UpperCAmelCase:					List[Any]		= min_tokens_to_keep







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase, __UpperCAmelCase:					List[Any]		= lax.top_k(snake_case_   ,    scores.shape[-1]					)

    __UpperCAmelCase:					Tuple		= jnp.full_like(snake_case_   ,    self.filter_value					)
    __UpperCAmelCase:					Optional[int]		= jax.nn.softmax(snake_case_   ,    axis=-1					).cumsum(axis=-1					)
    __UpperCAmelCase:					List[str]		= cumulative_probs < self.top_p

    # include the token that is higher than top_p as well
    __UpperCAmelCase:					str		= jnp.roll(snake_case_   ,    1					)
    score_mask |= score_mask.at[:, 0].set(snake_case_					)

    # min tokens to keep
    __UpperCAmelCase:					str		= score_mask.at[:, : self.min_tokens_to_keep].set(snake_case_					)

    __UpperCAmelCase:					List[Any]		= jnp.where(snake_case_   ,    snake_case_   ,    snake_case_					)
    __UpperCAmelCase:					Optional[Any]		= jax.lax.sort_key_val(snake_case_   ,    snake_case_					)[-1]

    return next_scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_ = -float("""Inf"""					)   ,    snake_case_ = 1					):

    '''simple docstring'''







    if not isinstance(snake_case_   ,    snake_case_					) or top_k <= 0:
      raise ValueError(F'''`top_k` has to be a strictly positive integer, but is {top_k}'''					)

    __UpperCAmelCase:					Optional[Any]		= max(snake_case_   ,    snake_case_					)
    __UpperCAmelCase:					List[str]		= filter_value







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase, __UpperCAmelCase:					List[Any]		= scores.shape
    __UpperCAmelCase:					List[Any]		= jnp.full(batch_size * vocab_size   ,    self.filter_value					)

    __UpperCAmelCase:					Any		= min(self.top_k   ,    scores.shape[-1]					)  # Safety check
    __UpperCAmelCase, __UpperCAmelCase:					List[Any]		= lax.top_k(snake_case_   ,    snake_case_					)
    __UpperCAmelCase:					Optional[Any]		= jnp.broadcast_to((jnp.arange(snake_case_					) * vocab_size)[:, None]   ,    (batch_size, topk)					).flatten()
    __UpperCAmelCase:					List[Any]		= topk_scores.flatten()
    __UpperCAmelCase:					str		= topk_indices.flatten() + shift

    __UpperCAmelCase:					List[Any]		= next_scores_flat.at[topk_indices_flat].set(snake_case_					)
    __UpperCAmelCase:					Any		= next_scores_flat.reshape(snake_case_   ,    snake_case_					)
    return next_scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					List[Any]		= bos_token_id







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					List[Any]		= jnp.full(scores.shape   ,    -float("""inf"""					)					)

    __UpperCAmelCase:					Optional[Any]		= 1 - jnp.bool_(cur_len - 1					)

    __UpperCAmelCase:					Any		= jnp.where(snake_case_   ,    new_scores.at[:, self.bos_token_id].set(0					)   ,    snake_case_					)

    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Union[str, Any]		= max_length
    __UpperCAmelCase:					List[str]		= eos_token_id







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[int]		= jnp.full(scores.shape   ,    -float("""inf"""					)					)

    __UpperCAmelCase:					List[str]		= 1 - jnp.bool_(cur_len - self.max_length + 1					)

    __UpperCAmelCase:					Any		= jnp.where(snake_case_   ,    new_scores.at[:, self.eos_token_id].set(0					)   ,    snake_case_					)

    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    if not isinstance(snake_case_   ,    snake_case_					) or min_length < 0:
      raise ValueError(F'''`min_length` has to be a positive integer, but is {min_length}'''					)

    if not isinstance(snake_case_   ,    snake_case_					) or eos_token_id < 0:
      raise ValueError(F'''`eos_token_id` has to be a positive integer, but is {eos_token_id}'''					)

    __UpperCAmelCase:					List[Any]		= min_length
    __UpperCAmelCase:					Optional[int]		= eos_token_id







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Union[str, Any]		= 1 - jnp.clip(cur_len - self.min_length   ,    0   ,    1					)

    __UpperCAmelCase:					int		= jnp.where(snake_case_   ,    scores.at[:, self.eos_token_id].set(-float("""inf"""					)					)   ,    snake_case_					)

    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Any		= list(snake_case_					)
    __UpperCAmelCase:					Optional[Any]		= begin_index







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Union[str, Any]		= 1 - jnp.bool_(cur_len - self.begin_index					)

    __UpperCAmelCase:					str		= jnp.where(snake_case_   ,    scores.at[:, self.begin_suppress_tokens].set(-float("""inf"""					)					)   ,    snake_case_					)

    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[Any]		= list(snake_case_					)







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					int		= scores.at[..., self.suppress_tokens].set(-float("""inf"""					)					)

    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[int]		= dict(snake_case_					)
    # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
    # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
    # Indexes without forced tokens will have a negative value.
    __UpperCAmelCase:					Optional[Any]		= jnp.ones((max(force_token_map.keys()					) + 1)   ,    dtype=jnp.intaa					) * -1
    for index, token in force_token_map.items():
      if token is not None:
        __UpperCAmelCase:					Union[str, Any]		= force_token_array.at[index].set(snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= jnp.intaa(snake_case_					)







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    def _force_token(snake_case_					):
      __UpperCAmelCase:					List[Any]		= scores.shape[0]
      __UpperCAmelCase:					int		= self.force_token_array[generation_idx]

      __UpperCAmelCase:					Union[str, Any]		= jnp.ones_like(snake_case_   ,    dtype=scores.dtype					) * -float("""inf"""					)
      __UpperCAmelCase:					List[Any]		= jnp.zeros((batch_size, 1)   ,    dtype=scores.dtype					)
      __UpperCAmelCase:					List[Any]		= lax.dynamic_update_slice(snake_case_   ,    snake_case_   ,    (0, current_token)					)
      return new_scores

    __UpperCAmelCase:					Optional[int]		= lax.cond(
        cur_len >= self.force_token_array.shape[0]   ,    lambda: scores   ,    lambda: lax.cond(
            self.force_token_array[cur_len] >= 0   ,    lambda: _force_token(snake_case_					)   ,    lambda: scores   ,    )   ,    )
    return scores




class    a (       __lowerCAmelCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					int		= generate_config.eos_token_id
    __UpperCAmelCase:					Union[str, Any]		= generate_config.no_timestamps_token_id
    __UpperCAmelCase:					Tuple		= generate_config.no_timestamps_token_id + 1

    __UpperCAmelCase:					Any		= decoder_input_length + 1

    if generate_config.is_multilingual:
      # room for language token and task token
      self.begin_index += 2
    if hasattr(snake_case_   ,    """max_initial_timestamp_index"""					):
      __UpperCAmelCase:					Optional[int]		= generate_config.max_initial_timestamp_index
    else:
      __UpperCAmelCase:					List[str]		= model_config.vocab_size
    if self.max_initial_timestamp_index is None:
      __UpperCAmelCase:					Any		= model_config.vocab_size







  def __call__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					Any		= scores.at[:, self.no_timestamps_token_id].set(-float("""inf"""					)					)

    def handle_pairs(snake_case_   ,    snake_case_					):
      __UpperCAmelCase:					Optional[int]		= jnp.where((cur_len - self.begin_index) >= 1   ,    snake_case_   ,    snake_case_					)
      __UpperCAmelCase:					Union[str, Any]		= jnp.where(
          input_ids_k[cur_len - 1] >= self.timestamp_begin   ,    True and last_was_timestamp   ,    snake_case_   ,    )

      __UpperCAmelCase:					Union[str, Any]		= jnp.where((cur_len - self.begin_index) < 2   ,    snake_case_   ,    snake_case_					)
      __UpperCAmelCase:					List[str]		= jnp.where(
          input_ids_k[cur_len - 2] >= self.timestamp_begin   ,    snake_case_   ,    snake_case_   ,    )

      return jnp.where(
          snake_case_   ,    jnp.where(
              penultimate_was_timestamp > 0   ,    scores_k.at[self.timestamp_begin :].set(-float("""inf"""					)					)   ,    scores_k.at[: self.eos_token_id].set(-float("""inf"""					)					)   ,    )   ,    snake_case_   ,    )

    __UpperCAmelCase:					Dict		= jax.vmap(snake_case_					)(snake_case_   ,    snake_case_					)

    __UpperCAmelCase:					Dict		= jnp.where(cur_len == self.begin_index   ,    snake_case_   ,    snake_case_					)
    __UpperCAmelCase:					List[Any]		= jnp.where(
        self.max_initial_timestamp_index is not None   ,    True and apply_max_initial_timestamp   ,    snake_case_   ,    )

    __UpperCAmelCase:					List[str]		= self.timestamp_begin + self.max_initial_timestamp_index

    __UpperCAmelCase:					List[str]		= jnp.where(
        snake_case_   ,    scores.at[:, last_allowed + 1 :].set(-float("""inf"""					)					)   ,    snake_case_   ,    )

    # if sum of probability over timestamps is above any other token, sample timestamp
    __UpperCAmelCase:					Union[str, Any]		= jax.nn.log_softmax(snake_case_   ,    axis=-1					)

    def handle_cumulative_probs(snake_case_   ,    snake_case_					):
      __UpperCAmelCase:					int		= jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]   ,    axis=-1					)
      __UpperCAmelCase:					Tuple		= jnp.max(logprobs_k[: self.timestamp_begin]					)
      return jnp.where(
          timestamp_logprob > max_text_token_logprob   ,    scores_k.at[: self.timestamp_begin].set(-float("""inf"""					)					)   ,    snake_case_   ,    )

    __UpperCAmelCase:					Tuple		= jax.vmap(snake_case_					)(snake_case_   ,    snake_case_					)

    return scores