



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 	_lowerCAmelCase						(				):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Any       =	10
 __SCREAMING_SNAKE_CASE       :    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"     ),
     }     )
 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	datasets.Dataset.from_dict(
     {
         "tokens": [["foo"] * 5] * n,
         "labels": [[1] * 5] * n,
         "answers": [{"answer_start": [97], "text": ["1976"]}] * 10,
         "id": list(range(__lowerCamelCase     )     ),
     }		,	features=__lowerCamelCase		,	)
 return dataset






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     Any     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Tuple       =	str(tmp_path_factory.mktemp("data"     ) / "file.arrow"     )
 dataset.map(cache_file_name=__lowerCamelCase     )
 return filename


# FILE_CONTENT + files


_lowerCamelCase						    =  """\
    Text data.
    Second line of data."""






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Union[str, Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	tmp_path_factory.mktemp("data"     ) / "file.txt"
 __SCREAMING_SNAKE_CASE       :    str       =	FILE_CONTENT
 with open(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase     )
 return filename






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 import bza

 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	tmp_path_factory.mktemp("data"     ) / "file.txt.bz2"
 __SCREAMING_SNAKE_CASE       :    str       =	bytes(__lowerCamelCase		,	"utf-8"     )
 with bza.open(__lowerCamelCase		,	"wb"     ) as f:
  f.write(__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 import gzip

 __SCREAMING_SNAKE_CASE       :    Dict       =	str(tmp_path_factory.mktemp("data"     ) / "file.txt.gz"     )
 __SCREAMING_SNAKE_CASE       :    int       =	bytes(__lowerCamelCase		,	"utf-8"     )
 with gzip.open(__lowerCamelCase		,	"wb"     ) as f:
  f.write(__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 if datasets.config.LZ4_AVAILABLE:
  import lza.frame

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






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     int		,	__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 if datasets.config.PY7ZR_AVAILABLE:
  import pyazr

  __SCREAMING_SNAKE_CASE       :    Any       =	tmp_path_factory.mktemp("data"     ) / "file.txt.7z"
  with pyazr.SevenZipFile(__lowerCamelCase		,	"w"     ) as archive:
   archive.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
  return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[Any]		,	__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 import tarfile

 __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	tmp_path_factory.mktemp("data"     ) / "file.txt.tar"
 with tarfile.TarFile(__lowerCamelCase		,	"w"     ) as f:
  f.add(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     str     ):



 """simple docstring"""





 import lzma

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






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Union[str, Any]		,	__lowerCamelCase   :     Optional[Any]     ):



 """simple docstring"""





 import zipfile

 __SCREAMING_SNAKE_CASE       :    List[Any]       =	tmp_path_factory.mktemp("data"     ) / "file.txt.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Union[str, Any]     ):



 """simple docstring"""





 if datasets.config.ZSTANDARD_AVAILABLE:
  import zstandard as zstd

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






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Dict       =	tmp_path_factory.mktemp("data"     ) / "file.xml"
 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	textwrap.dedent(
     "\\n    <?xml version=\"1.0\" encoding=\"UTF-8\" ?>\n    <tmx version=\"1.4\">\n      <header segtype=\"sentence\" srclang=\"ca\" />\n      <body>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 1</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 1</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 2</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 2</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 3</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 3</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 4</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 4</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 5</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 5</seg></tuv>\n        </tu>\n      </body>\n    </tmx>"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase     )
 return filename


_lowerCamelCase						    =  [
    {"""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},
]
_lowerCamelCase						    =  [
    {"""col_1""": """4""", """col_2""": 4, """col_3""": 4.0},
    {"""col_1""": """5""", """col_2""": 5, """col_3""": 5.0},
]
_lowerCamelCase						    =  {
    """col_1""": ["""0""", """1""", """2""", """3"""],
    """col_2""": [0, 1, 2, 3],
    """col_3""": [0.0, 1.0, 2.0, 3.0],
}

_lowerCamelCase						    =  [
    {"""col_3""": 0.0, """col_1""": """0""", """col_2""": 0},
    {"""col_3""": 1.0, """col_1""": """1""", """col_2""": 1},
]

_lowerCamelCase						    =  [
    {"""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 	_lowerCAmelCase						(				):



 """simple docstring"""





 return DATA_DICT_OF_LISTS






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[int]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[str]       =	datasets.Dataset.from_dict(__lowerCamelCase     )
 __SCREAMING_SNAKE_CASE       :    str       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.arrow"     )
 dataset.map(cache_file_name=__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     str     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.sqlite"     )
 with contextlib.closing(sqlitea.connect(__lowerCamelCase     )     ) as con:
  __SCREAMING_SNAKE_CASE       :    Tuple       =	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 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.csv"     )
 with open(__lowerCamelCase		,	"w"		,	newline=""     ) as f:
  __SCREAMING_SNAKE_CASE       :    List[Any]       =	csv.DictWriter(__lowerCamelCase		,	fieldnames=["col_1", "col_2", "col_3"]     )
  writer.writeheader()
  for item in DATA:
   writer.writerow(__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	str(tmp_path_factory.mktemp("data"     ) / "dataset2.csv"     )
 with open(__lowerCamelCase		,	"w"		,	newline=""     ) as f:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	csv.DictWriter(__lowerCamelCase		,	fieldnames=["col_1", "col_2", "col_3"]     )
  writer.writeheader()
  for item in DATA:
   writer.writerow(__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Union[str, Any]		,	__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 import bza

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






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[int]		,	__lowerCamelCase   :     Tuple		,	__lowerCamelCase   :     Tuple     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	tmp_path_factory.mktemp("data"     ) / "dataset.csv.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     Any		,	__lowerCamelCase   :     Any     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[Any]       =	tmp_path_factory.mktemp("data"     ) / "dataset.csv.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename(csv_path.replace(".csv"		,	".CSV"     )     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.basename(csva_path.replace(".csv"		,	".CSV"     )     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[Any]		,	__lowerCamelCase   :     Any		,	__lowerCamelCase   :     str     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Tuple       =	tmp_path_factory.mktemp("data"     ) / "dataset_with_dir.csv.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.join("main_dir"		,	os.path.basename(__lowerCamelCase     )     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.join("main_dir"		,	os.path.basename(__lowerCamelCase     )     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.parquet"     )
 __SCREAMING_SNAKE_CASE       :    int       =	pa.schema(
     {
         "col_1": pa.string(),
         "col_2": pa.intaa(),
         "col_3": pa.floataa(),
     }     )
 with open(__lowerCamelCase		,	"wb"     ) as f:
  __SCREAMING_SNAKE_CASE       :    Dict       =	pq.ParquetWriter(__lowerCamelCase		,	schema=__lowerCamelCase     )
  __SCREAMING_SNAKE_CASE       :    int       =	pa.Table.from_pydict({k: [DATA[i][k] for i in range(len(__lowerCamelCase     )     )] for k in DATA[0]}		,	schema=__lowerCamelCase     )
  writer.write_table(__lowerCamelCase     )
  writer.close()
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.json"     )
 __SCREAMING_SNAKE_CASE       :    int       =	{"data": DATA}
 with open(__lowerCamelCase		,	"w"     ) as f:
  json.dump(__lowerCamelCase		,	__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Union[str, Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.json"     )
 __SCREAMING_SNAKE_CASE       :    Dict       =	{"data": DATA_DICT_OF_LISTS}
 with open(__lowerCamelCase		,	"w"     ) as f:
  json.dump(__lowerCamelCase		,	__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     str     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[str]       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.jsonl"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in DATA:
   f.write(json.dumps(__lowerCamelCase     ) + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Any       =	str(tmp_path_factory.mktemp("data"     ) / "dataset2.jsonl"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in DATA:
   f.write(json.dumps(__lowerCamelCase     ) + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Any       =	str(tmp_path_factory.mktemp("data"     ) / "dataset_312.jsonl"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in DATA_312:
   f.write(json.dumps(__lowerCamelCase     ) + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Any       =	str(tmp_path_factory.mktemp("data"     ) / "dataset-str.jsonl"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in DATA_STR:
   f.write(json.dumps(__lowerCamelCase     ) + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     int		,	__lowerCamelCase   :     Any     ):



 """simple docstring"""





 import gzip

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






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     int		,	__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 import gzip

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






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     Optional[Any]		,	__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	tmp_path_factory.mktemp("data"     ) / "dataset.jsonl.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Tuple		,	__lowerCamelCase   :     Tuple		,	__lowerCamelCase   :     Optional[Any]		,	__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[str]       =	tmp_path_factory.mktemp("data"     ) / "dataset_nested.jsonl.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.join("nested"		,	os.path.basename(__lowerCamelCase     )     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[int]		,	__lowerCamelCase   :     int		,	__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	tmp_path_factory.mktemp("data"     ) / "dataset_with_dir.jsonl.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.join("main_dir"		,	os.path.basename(__lowerCamelCase     )     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.join("main_dir"		,	os.path.basename(__lowerCamelCase     )     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     int		,	__lowerCamelCase   :     Optional[int]		,	__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	tmp_path_factory.mktemp("data"     ) / "dataset.jsonl.tar"
 with tarfile.TarFile(__lowerCamelCase		,	"w"     ) as f:
  f.add(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
  f.add(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     str		,	__lowerCamelCase   :     Any		,	__lowerCamelCase   :     str		,	__lowerCamelCase   :     Tuple     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[Any]       =	tmp_path_factory.mktemp("data"     ) / "dataset_nested.jsonl.tar"
 with tarfile.TarFile(__lowerCamelCase		,	"w"     ) as f:
  f.add(__lowerCamelCase		,	arcname=os.path.join("nested"		,	os.path.basename(__lowerCamelCase     )     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	["0", "1", "2", "3"]
 __SCREAMING_SNAKE_CASE       :    Dict       =	str(tmp_path_factory.mktemp("data"     ) / "dataset.txt"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in data:
   f.write(item + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	["0", "1", "2", "3"]
 __SCREAMING_SNAKE_CASE       :    List[str]       =	str(tmp_path_factory.mktemp("data"     ) / "dataset2.txt"     )
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in data:
   f.write(item + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[Any]       =	["0", "1", "2", "3"]
 __SCREAMING_SNAKE_CASE       :    Any       =	tmp_path_factory.mktemp("data"     ) / "dataset.abc"
 with open(__lowerCamelCase		,	"w"     ) as f:
  for item in data:
   f.write(item + "\n"     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     int		,	__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     str     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	tmp_path_factory.mktemp("data"     ) / "dataset.text.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[Any]		,	__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     Tuple     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[str]       =	tmp_path_factory.mktemp("data"     ) / "dataset_with_dir.text.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.join("main_dir"		,	os.path.basename(__lowerCamelCase     )     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.join("main_dir"		,	os.path.basename(__lowerCamelCase     )     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	tmp_path_factory.mktemp("data"     ) / "dataset.ext.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename("unsupported.ext"     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.basename("unsupported_2.ext"     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    str       =	"\n".join(["First", "Second\u2029with Unicode new line", "Third"]     )
 __SCREAMING_SNAKE_CASE       :    str       =	str(tmp_path_factory.mktemp("data"     ) / "dataset_with_unicode_new_lines.txt"     )
 with open(__lowerCamelCase		,	"w"		,	encoding="utf-8"     ) as f:
  f.write(__lowerCamelCase     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				):



 """simple docstring"""





 return os.path.join("tests"		,	"features"		,	"data"		,	"test_image_rgb.jpg"     )






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				):



 """simple docstring"""





 return os.path.join("tests"		,	"features"		,	"data"		,	"test_audio_44100.wav"     )






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    str       =	tmp_path_factory.mktemp("data"     ) / "dataset.img.zip"
 with zipfile.ZipFile(__lowerCamelCase		,	"w"     ) as f:
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     )     )
  f.write(__lowerCamelCase		,	arcname=os.path.basename(__lowerCamelCase     ).replace(".jpg"		,	"2.jpg"     )     )
 return path






@pytest.fixture(scope="session"     )
def 	_lowerCAmelCase						(				__lowerCamelCase   :     str     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	tmp_path_factory.mktemp("data_dir"     )

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

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

 return data_dir











import argparse
import json
import os
from pathlib import Path

import requests
import torch

from transformers import JukeboxConfig, JukeboxModel
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase						    =  logging.get_logger(__name__)


_lowerCamelCase						    =  """https://openaipublic.azureedge.net/jukebox/models/"""
_lowerCamelCase						    =  {
    """jukebox-1b-lyrics""": [
        """5b/vqvae.pth.tar""",
        """5b/prior_level_0.pth.tar""",
        """5b/prior_level_1.pth.tar""",
        """1b_lyrics/prior_level_2.pth.tar""",
    ],
    """jukebox-5b-lyrics""": [
        """5b/vqvae.pth.tar""",
        """5b/prior_level_0.pth.tar""",
        """5b/prior_level_1.pth.tar""",
        """5b_lyrics/prior_level_2.pth.tar""",
    ],
}






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 if key.endswith(".model.1.bias"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    Dict       =	key.replace(".model.1.bias"		,	".conv1d_1.bias"     )
 elif key.endswith(".model.1.weight"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    Tuple       =	key.replace(".model.1.weight"		,	".conv1d_1.weight"     )
 elif key.endswith(".model.3.bias"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    Dict       =	key.replace(".model.3.bias"		,	".conv1d_2.bias"     )
 elif key.endswith(".model.3.weight"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	key.replace(".model.3.weight"		,	".conv1d_2.weight"     )

 if "conditioner_blocks.0." in key:
  __SCREAMING_SNAKE_CASE       :    List[Any]       =	key.replace("conditioner_blocks.0"		,	"conditioner_blocks"     )

 if "prime_prior" in key:
  __SCREAMING_SNAKE_CASE       :    Dict       =	key.replace("prime_prior"		,	"encoder"     )

 if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
  __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	key.replace(".emb."		,	"."     )

 if key.endswith("k"     ):  # replace vqvae.X.k with vqvae.X.codebook
  return key.replace(".k"		,	".codebook"     )
 if "y_emb." in key:
  return key.replace("y_emb."		,	"metadata_embedding."     )

 if "x_emb.emb." in key:
  __SCREAMING_SNAKE_CASE       :    Any       =	key.replace("0.x_emb.emb"		,	"embed_tokens"     )

 if "prime_state_ln" in key:
  return key.replace("prime_state_ln"		,	"encoder.final_layer_norm"     )
 if ".ln" in key:
  return key.replace(".ln"		,	".layer_norm"     )
 if "_ln" in key:
  return key.replace("_ln"		,	"_layer_norm"     )

 if "prime_state_proj" in key:
  return key.replace("prime_state_proj"		,	"encoder.proj_in"     )
 if "prime_x_out" in key:
  return key.replace("prime_x_out"		,	"encoder.lm_head"     )
 if "prior.x_out" in key:
  return key.replace("x_out"		,	"fc_proj_out"     )
 if "x_emb" in key:
  return key.replace("x_emb"		,	"embed_tokens"     )

 return key






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Tuple		,	__lowerCamelCase   :     Union[str, Any]		,	__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     Optional[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    str       =	{}
 import re

 __SCREAMING_SNAKE_CASE       :    str       =	re.compile(r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Tuple       =	re.compile(
     r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Tuple       =	re.compile(r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)"     )

 __SCREAMING_SNAKE_CASE       :    int       =	re.compile(r"decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re.compile(
     r"decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    List[str]       =	re.compile(r"decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)"     )

 __SCREAMING_SNAKE_CASE       :    int       =	re.compile(r"conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    List[str]       =	re.compile(
     r"conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Dict       =	re.compile(r"conditioner_blocks.(\d*).cond.model.(\d*).(bias|weight)"     )

 for original_key, value in state_dict.items():
  # rename vqvae.encoder keys
  if re_encoder_block_conv_in.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	re_encoder_block_conv_in.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    List[str]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	int(groups[2]     ) * 2 + int(groups[3]     )
   __SCREAMING_SNAKE_CASE       :    int       =	F"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_encoder_block_conv_in.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_encoder_block_resnet.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Dict       =	re_encoder_block_resnet.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    List[str]       =	int(groups[2]     ) * 2 + int(groups[3]     )
   __SCREAMING_SNAKE_CASE       :    Any       =	{"1": 1, "3": 2}[groups[-2]]
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	F"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}."""
   __SCREAMING_SNAKE_CASE       :    List[str]       =	F"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Any       =	prefix + resnet_block
   __SCREAMING_SNAKE_CASE       :    List[str]       =	re_encoder_block_resnet.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_encoder_block_proj_out.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Any       =	re_encoder_block_proj_out.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    int       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	F"""encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Dict       =	re_encoder_block_proj_out.sub(__lowerCamelCase		,	__lowerCamelCase     )

  # rename vqvae.decoder keys
  elif re_decoder_block_conv_out.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    List[str]       =	re_decoder_block_conv_out.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    str       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	int(groups[2]     ) * 2 + int(groups[3]     ) - 2
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	F"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	re_decoder_block_conv_out.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_decoder_block_resnet.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_decoder_block_resnet.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    int       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	int(groups[2]     ) * 2 + int(groups[3]     ) - 2
   __SCREAMING_SNAKE_CASE       :    Any       =	{"1": 1, "3": 2}[groups[-2]]
   __SCREAMING_SNAKE_CASE       :    Tuple       =	F"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}."""
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	F"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	prefix + resnet_block
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_decoder_block_resnet.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_decoder_block_proj_in.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	re_decoder_block_proj_in.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    List[str]       =	F"""decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    int       =	re_decoder_block_proj_in.sub(__lowerCamelCase		,	__lowerCamelCase     )

  # rename prior cond.model to upsampler.upsample_block and resnet
  elif re_prior_cond_conv_out.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	re_prior_cond_conv_out.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    str       =	int(groups[1]     ) * 2 + int(groups[2]     ) - 2
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	F"""conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Any       =	re_prior_cond_conv_out.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_prior_cond_resnet.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_prior_cond_resnet.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Any       =	int(groups[1]     ) * 2 + int(groups[2]     ) - 2
   __SCREAMING_SNAKE_CASE       :    str       =	{"1": 1, "3": 2}[groups[-2]]
   __SCREAMING_SNAKE_CASE       :    Any       =	F"""conditioner_blocks.upsampler.upsample_block.{block_index}."""
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	F"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Any       =	prefix + resnet_block
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_prior_cond_resnet.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_prior_cond_proj_in.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Tuple       =	re_prior_cond_proj_in.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Dict       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Dict       =	F"""conditioner_blocks.upsampler.proj_in.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    int       =	re_prior_cond_proj_in.sub(__lowerCamelCase		,	__lowerCamelCase     )

  # keep original key
  else:
   __SCREAMING_SNAKE_CASE       :    Tuple       =	original_key

  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	replace_key(__lowerCamelCase     )

  if F"""{key_prefix}.{key}""" not in model_state_dict or key is None:
   print(F"""failed converting {original_key} to {key}, does not match"""     )

  # handle missmatched shape
  elif value.shape != model_state_dict[F"""{key_prefix}.{key}"""].shape:
   __SCREAMING_SNAKE_CASE       :    List[str]       =	model_state_dict[F"""{key_prefix}.{key}"""]
   print(F"""{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match"""     )
   __SCREAMING_SNAKE_CASE       :    str       =	original_key

  __SCREAMING_SNAKE_CASE       :    List[str]       =	original_key
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	value

 return new_dict






@torch.no_grad()
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[int]=None		,	__lowerCamelCase   :     List[Any]=None     ):



 """simple docstring"""





 for file in MODEL_MAPPING[model_name]:
  if not os.path.isfile(F"""{pytorch_dump_folder_path}/{file.split("/"     )[-1]}"""     ):
   __SCREAMING_SNAKE_CASE       :    Dict       =	requests.get(F"""{PREFIX}{file}"""		,	allow_redirects=__lowerCamelCase     )
   os.makedirs(F"""{pytorch_dump_folder_path}/"""		,	exist_ok=__lowerCamelCase     )
   open(F"""{pytorch_dump_folder_path}/{file.split("/"     )[-1]}"""		,	"wb"     ).write(r.content     )

 __SCREAMING_SNAKE_CASE       :    int       =	MODEL_MAPPING[model_name.split("/"     )[-1]]

 __SCREAMING_SNAKE_CASE       :    List[str]       =	JukeboxConfig.from_pretrained(__lowerCamelCase     )
 __SCREAMING_SNAKE_CASE       :    str       =	JukeboxModel(__lowerCamelCase     )

 __SCREAMING_SNAKE_CASE       :    int       =	[]
 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	{}
 for i, dict_name in enumerate(__lowerCamelCase     ):
  __SCREAMING_SNAKE_CASE       :    str       =	torch.load(F"""{pytorch_dump_folder_path}/{dict_name.split("/"     )[-1]}"""     )["model"]

  __SCREAMING_SNAKE_CASE       :    Optional[int]       =	{}
  for k in old_dic.keys():
   if k.endswith(".b"     ):
    __SCREAMING_SNAKE_CASE       :    Optional[int]       =	old_dic[k]
   elif k.endswith(".w"     ):
    __SCREAMING_SNAKE_CASE       :    int       =	old_dic[k]
   elif "level_2" not in dict_name and "cond.model." in k:
    __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	old_dic[k]
   else:
    __SCREAMING_SNAKE_CASE       :    Optional[int]       =	old_dic[k]

  __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	"vqvae" if i == 0 else F"""priors.{3 - i}"""
  __SCREAMING_SNAKE_CASE       :    int       =	fix_jukebox_keys(__lowerCamelCase		,	model.state_dict()		,	__lowerCamelCase		,	__lowerCamelCase     )
  weight_dict.append(__lowerCamelCase     )

 __SCREAMING_SNAKE_CASE       :    Dict       =	weight_dict.pop(0     )
 model.vqvae.load_state_dict(__lowerCamelCase     )
 for i in range(len(__lowerCamelCase     )     ):
  model.priors[i].load_state_dict(weight_dict[2 - i]     )

 Path(__lowerCamelCase     ).mkdir(exist_ok=__lowerCamelCase     )
 with open(F"""{pytorch_dump_folder_path}/mapping.json"""		,	"w"     ) as txtfile:
  json.dump(__lowerCamelCase		,	__lowerCamelCase     )

 print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""     )
 model.save_pretrained(__lowerCamelCase     )

 return weight_dict


if __name__ == "__main__":
      _lowerCamelCase						    =  argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--model_name""",
          default="""jukebox-5b-lyrics""",
          type=str,
          help="""Name of the model you'd like to convert.""",
      )
      parser.add_argument(
          """--pytorch_dump_folder_path""",
          default="""jukebox-5b-lyrics-converted""",
          type=str,
          help="""Path to the output PyTorch model directory.""",
      )
      _lowerCamelCase						    =  parser.parse_args()
      convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)






