


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      snake_case_						( ):
   UpperCAmelCase_							:      int						=   1_0
   UpperCAmelCase_							:      Tuple						=   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_							:      Dict						=   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      snake_case_						( __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Any						=   str(tmp_path_factory.mktemp('''data'''       ) / '''file.arrow'''       )
   dataset.map(cache_file_name=__lowercase       )
   return filename


# FILE_CONTENT + files


__UpperCamelCase				:  Optional[int]      					=							'\\n    Text data.\n    Second line of data.'
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Dict						=   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      snake_case_						( __lowercase       ):
   import bza

   UpperCAmelCase_							:      List[str]						=   tmp_path_factory.mktemp('''data'''       ) / '''file.txt.bz2'''
   UpperCAmelCase_							:      Dict						=   bytes(__lowercase			,       '''utf-8'''       )
   with bza.open(__lowercase			,       '''wb'''       ) as f:
      f.write(__lowercase       )
   return path
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase       ):
   import gzip

   UpperCAmelCase_							:      Union[str, 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      snake_case_						( __lowercase       ):
   if datasets.config.LZ4_AVAILABLE:
      import lza.frame

      UpperCAmelCase_							:      Optional[Any]						=   tmp_path_factory.mktemp('''data'''       ) / '''file.txt.lz4'''
      UpperCAmelCase_							:      Tuple						=   bytes(__lowercase			,       '''utf-8'''       )
      with lza.frame.open(__lowercase			,       '''wb'''       ) as f:
         f.write(__lowercase       )
      return path
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase			,       __lowercase       ):
   if datasets.config.PY7ZR_AVAILABLE:
      import pyazr

      UpperCAmelCase_							:      Dict						=   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      snake_case_						( __lowercase			,       __lowercase       ):
   import tarfile

   UpperCAmelCase_							:      Any						=   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      snake_case_						( __lowercase       ):
   import lzma

   UpperCAmelCase_							:      List[Any]						=   tmp_path_factory.mktemp('''data'''       ) / '''file.txt.xz'''
   UpperCAmelCase_							:      List[str]						=   bytes(__lowercase			,       '''utf-8'''       )
   with lzma.open(__lowercase			,       '''wb'''       ) as f:
      f.write(__lowercase       )
   return path
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase			,       __lowercase       ):
   import zipfile

   UpperCAmelCase_							:      Any						=   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      snake_case_						( __lowercase       ):
   if datasets.config.ZSTANDARD_AVAILABLE:
      import zstandard as zstd

      UpperCAmelCase_							:      Tuple						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Optional[int]						=   tmp_path_factory.mktemp('''data'''       ) / '''file.xml'''
   UpperCAmelCase_							:      Tuple						=   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


__UpperCamelCase				:  Optional[int]      					=							[
    {'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},
]
__UpperCamelCase				:  Union[str, Any]      					=							[
    {'col_1': '4', 'col_2': 4, 'col_3': 4.0},
    {'col_1': '5', 'col_2': 5, 'col_3': 5.0},
]
__UpperCamelCase				:  Union[str, Any]      					=							{
    'col_1': ['0', '1', '2', '3'],
    'col_2': [0, 1, 2, 3],
    'col_3': [0.0, 1.0, 2.0, 3.0],
}

__UpperCamelCase				:  Optional[Any]      					=							[
    {'col_3': 0.0, 'col_1': '0', 'col_2': 0},
    {'col_3': 1.0, 'col_1': '1', 'col_2': 1},
]

__UpperCamelCase				:  Tuple      					=							[
    {'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      snake_case_						( ):
   return DATA_DICT_OF_LISTS
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Dict						=   datasets.Dataset.from_dict(__lowercase       )
   UpperCAmelCase_							:      Union[str, Any]						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset.arrow'''       )
   dataset.map(cache_file_name=__lowercase       )
   return path
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      int						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset.sqlite'''       )
   with contextlib.closing(sqlitea.connect(__lowercase       )       ) as con:
      UpperCAmelCase_							:      int						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset.csv'''       )
   with open(__lowercase			,       '''w'''			,       newline=''''''       ) as f:
      UpperCAmelCase_							:      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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Tuple						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset2.csv'''       )
   with open(__lowercase			,       '''w'''			,       newline=''''''       ) as f:
      UpperCAmelCase_							:      Union[str, 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      snake_case_						( __lowercase			,       __lowercase       ):
   import bza

   UpperCAmelCase_							:      int						=   tmp_path_factory.mktemp('''data'''       ) / '''dataset.csv.bz2'''
   with open(__lowercase			,       '''rb'''       ) as f:
      UpperCAmelCase_							:      Optional[int]						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      List[str]						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Any						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      str						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset.parquet'''       )
   UpperCAmelCase_							:      Tuple						=   pa.schema(
       {
           '''col_1''': pa.string(),
           '''col_2''': pa.intaa(),
           '''col_3''': pa.floataa(),
       }       )
   with open(__lowercase			,       '''wb'''       ) as f:
      UpperCAmelCase_							:      Union[str, Any]						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Optional[Any]						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset.json'''       )
   UpperCAmelCase_							:      Optional[Any]						=   {'''data''': DATA}
   with open(__lowercase			,       '''w'''       ) as f:
      json.dump(__lowercase			,       __lowercase       )
   return path
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Optional[int]						=   str(tmp_path_factory.mktemp('''data'''       ) / '''dataset.json'''       )
   UpperCAmelCase_							:      int						=   {'''data''': DATA_DICT_OF_LISTS}
   with open(__lowercase			,       '''w'''       ) as f:
      json.dump(__lowercase			,       __lowercase       )
   return path
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Union[str, 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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      str						=   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      snake_case_						( __lowercase			,       __lowercase       ):
   import gzip

   UpperCAmelCase_							:      Optional[int]						=   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      snake_case_						( __lowercase			,       __lowercase       ):
   import gzip

   UpperCAmelCase_							:      Dict						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Optional[int]						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Dict						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Optional[Any]						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase			,       __lowercase       ):
   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Tuple						=   ['''0''', '''1''', '''2''', '''3''']
   UpperCAmelCase_							:      Dict						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Union[str, Any]						=   ['''0''', '''1''', '''2''', '''3''']
   UpperCAmelCase_							:      List[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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      Dict						=   ['''0''', '''1''', '''2''', '''3''']
   UpperCAmelCase_							:      str						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Any						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Any						=   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      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      str						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      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      snake_case_						( ):
   return os.path.join('''tests'''			,       '''features'''			,       '''data'''			,       '''test_image_rgb.jpg'''       )
@pytest.fixture(scope='''session'''       )
def      snake_case_						( ):
   return os.path.join('''tests'''			,       '''features'''			,       '''data'''			,       '''test_audio_44100.wav'''       )
@pytest.fixture(scope='''session'''       )
def      snake_case_						( __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      List[Any]						=   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      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      int						=   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



import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
def      snake_case_						( ):
   UpperCAmelCase_							:      int						=   '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
   UpperCAmelCase_							:      Dict						=   Image.open(requests.get(__lowercase			,       stream=__lowercase       ).raw       ).convert('''RGB'''       )

   return image
def      snake_case_						( __lowercase       ):
   UpperCAmelCase_							:      List[str]						=   []
   # fmt: off

   # vision encoder
   rename_keys.append(('''visual_encoder.cls_token''', '''vision_model.embeddings.class_embedding''')       )
   rename_keys.append(('''visual_encoder.pos_embed''', '''vision_model.embeddings.position_embedding''')       )
   rename_keys.append(('''visual_encoder.patch_embed.proj.weight''', '''vision_model.embeddings.patch_embedding.weight''')       )
   rename_keys.append(('''visual_encoder.patch_embed.proj.bias''', '''vision_model.embeddings.patch_embedding.bias''')       )
   rename_keys.append(('''ln_vision.weight''', '''vision_model.post_layernorm.weight''')       )
   rename_keys.append(('''ln_vision.bias''', '''vision_model.post_layernorm.bias''')       )

   for i in range(config.vision_config.num_hidden_layers       ):
      rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.weight''', F'''vision_model.encoder.layers.{i}.layer_norm1.weight''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.norm1.bias''', F'''vision_model.encoder.layers.{i}.layer_norm1.bias''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.weight''', F'''vision_model.encoder.layers.{i}.layer_norm2.weight''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.norm2.bias''', F'''vision_model.encoder.layers.{i}.layer_norm2.bias''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.attn.qkv.weight''', F'''vision_model.encoder.layers.{i}.self_attn.qkv.weight''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.weight''', F'''vision_model.encoder.layers.{i}.self_attn.projection.weight''',)       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.attn.proj.bias''', F'''vision_model.encoder.layers.{i}.self_attn.projection.bias''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc1.weight''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc1.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc1.bias''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.weight''', F'''vision_model.encoder.layers.{i}.mlp.fc2.weight''')       )
      rename_keys.append((F'''visual_encoder.blocks.{i}.mlp.fc2.bias''', F'''vision_model.encoder.layers.{i}.mlp.fc2.bias''')       )

   # QFormer
   rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.weight''', '''qformer.layernorm.weight''')       )
   rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.layernorm.bias''')       )

   # fmt: on
   return rename_keys
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Any						=   dct.pop(__lowercase       )
   UpperCAmelCase_							:      Optional[Any]						=   val
def      snake_case_						( __lowercase			,       __lowercase       ):
   for i in range(config.vision_config.num_hidden_layers       ):
      # read in original q and v biases
      UpperCAmelCase_							:      Optional[int]						=   state_dict.pop(F'''visual_encoder.blocks.{i}.attn.q_bias'''       )
      UpperCAmelCase_							:      Any						=   state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''       )

      # next, set bias in the state dict
      UpperCAmelCase_							:      int						=   torch.cat((q_bias, torch.zeros_like(__lowercase			,       requires_grad=__lowercase       ), v_bias)       )
      UpperCAmelCase_							:      List[str]						=   qkv_bias
def      snake_case_						( __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Optional[Any]						=   3_6_4 if '''coco''' in model_name else 2_2_4
   UpperCAmelCase_							:      Any						=   BlipaVisionConfig(image_size=__lowercase       ).to_dict()

   # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
   # seems like flan-T5 models don't have bos_token_id properly set?
   if "opt-2.7b" in model_name:
      UpperCAmelCase_							:      Any						=   OPTConfig.from_pretrained('''facebook/opt-2.7b'''			,       eos_token_id=__lowercase       ).to_dict()
   elif "opt-6.7b" in model_name:
      UpperCAmelCase_							:      List[str]						=   OPTConfig.from_pretrained('''facebook/opt-6.7b'''			,       eos_token_id=__lowercase       ).to_dict()
   elif "t5-xl" in model_name:
      UpperCAmelCase_							:      List[str]						=   TaConfig.from_pretrained('''google/flan-t5-xl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1       ).to_dict()
   elif "t5-xxl" in model_name:
      UpperCAmelCase_							:      Any						=   TaConfig.from_pretrained('''google/flan-t5-xxl'''			,       dense_act_fn='''gelu'''			,       bos_token_id=1       ).to_dict()

   UpperCAmelCase_							:      List[Any]						=   BlipaConfig(vision_config=__lowercase			,       text_config=__lowercase       )

   return config, image_size
@torch.no_grad()
def      snake_case_						( __lowercase			,       __lowercase=None			,       __lowercase=False       ):
   UpperCAmelCase_							:      List[Any]						=   (
       AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''       )
       if '''opt''' in model_name
       else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''       )
   )
   UpperCAmelCase_							:      List[str]						=   tokenizer('''\n'''			,       add_special_tokens=__lowercase       ).input_ids[0]
   UpperCAmelCase_		,      UpperCAmelCase_							:      str						=   get_blipa_config(__lowercase			,       eos_token_id=__lowercase       )

   UpperCAmelCase_							:      List[Any]						=   BlipaForConditionalGeneration(__lowercase       ).eval()

   UpperCAmelCase_							:      Tuple						=   {
       '''blip2-opt-2.7b''': ('''blip2_opt''', '''pretrain_opt2.7b'''),
       '''blip2-opt-6.7b''': ('''blip2_opt''', '''pretrain_opt6.7b'''),
       '''blip2-opt-2.7b-coco''': ('''blip2_opt''', '''caption_coco_opt2.7b'''),
       '''blip2-opt-6.7b-coco''': ('''blip2_opt''', '''caption_coco_opt6.7b'''),
       '''blip2-flan-t5-xl''': ('''blip2_t5''', '''pretrain_flant5xl'''),
       '''blip2-flan-t5-xl-coco''': ('''blip2_t5''', '''caption_coco_flant5xl'''),
       '''blip2-flan-t5-xxl''': ('''blip2_t5''', '''pretrain_flant5xxl'''),
   }

   UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[Any]						=   model_name_to_original[model_name]

   # load original model
   print('''Loading original model...'''       )
   UpperCAmelCase_							:      int						=   '''cuda''' if torch.cuda.is_available() else '''cpu'''
   UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[Any]						=   load_model_and_preprocess(
       name=__lowercase			,       model_type=__lowercase			,       is_eval=__lowercase			,       device=__lowercase       )
   original_model.eval()
   print('''Done!'''       )

   # update state dict keys
   UpperCAmelCase_							:      Optional[Any]						=   original_model.state_dict()
   UpperCAmelCase_							:      List[Any]						=   create_rename_keys(__lowercase       )
   for src, dest in rename_keys:
      rename_key(__lowercase			,       __lowercase			,       __lowercase       )

   # some keys can be renamed efficiently
   for key, val in state_dict.copy().items():
      UpperCAmelCase_							:      Union[str, Any]						=   state_dict.pop(__lowercase       )
      if key.startswith('''Qformer.bert'''       ):
         UpperCAmelCase_							:      Tuple						=   key.replace('''Qformer.bert'''			,       '''qformer'''       )
      if "attention.self" in key:
         UpperCAmelCase_							:      Optional[Any]						=   key.replace('''self'''			,       '''attention'''       )
      if "opt_proj" in key:
         UpperCAmelCase_							:      Any						=   key.replace('''opt_proj'''			,       '''language_projection'''       )
      if "t5_proj" in key:
         UpperCAmelCase_							:      Tuple						=   key.replace('''t5_proj'''			,       '''language_projection'''       )
      if key.startswith('''opt'''       ):
         UpperCAmelCase_							:      Any						=   key.replace('''opt'''			,       '''language'''       )
      if key.startswith('''t5'''       ):
         UpperCAmelCase_							:      Optional[Any]						=   key.replace('''t5'''			,       '''language'''       )
      UpperCAmelCase_							:      List[str]						=   val

   # read in qv biases
   read_in_q_v_bias(__lowercase			,       __lowercase       )

   UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[int]						=   hf_model.load_state_dict(__lowercase			,       strict=__lowercase       )
   assert len(__lowercase       ) == 0
   assert unexpected_keys == ["qformer.embeddings.position_ids"]

   UpperCAmelCase_							:      str						=   load_demo_image()
   UpperCAmelCase_							:      Any						=   vis_processors['''eval'''](__lowercase       ).unsqueeze(0       ).to(__lowercase       )
   UpperCAmelCase_							:      Optional[Any]						=   tokenizer(['''\n''']			,       return_tensors='''pt'''       ).input_ids.to(__lowercase       )

   # create processor
   UpperCAmelCase_							:      Optional[int]						=   BlipImageProcessor(
       size={'''height''': image_size, '''width''': image_size}			,       image_mean=__lowercase			,       image_std=__lowercase       )
   UpperCAmelCase_							:      Tuple						=   BlipaProcessor(image_processor=__lowercase			,       tokenizer=__lowercase       )
   UpperCAmelCase_							:      str						=   processor(images=__lowercase			,       return_tensors='''pt'''       ).pixel_values.to(__lowercase       )

   # make sure processor creates exact same pixel values
   assert torch.allclose(__lowercase			,       __lowercase       )

   original_model.to(__lowercase       )
   hf_model.to(__lowercase       )
   with torch.no_grad():
      if "opt" in model_name:
         UpperCAmelCase_							:      Tuple						=   original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}       ).logits
         UpperCAmelCase_							:      Optional[int]						=   hf_model(__lowercase			,       __lowercase       ).logits
      else:
         UpperCAmelCase_							:      int						=   original_model(
             {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}       ).logits
         UpperCAmelCase_							:      Optional[int]						=   input_ids.masked_fill(input_ids == tokenizer.pad_token_id			,       -1_0_0       )
         UpperCAmelCase_							:      int						=   hf_model(__lowercase			,       __lowercase			,       labels=__lowercase       ).logits

   assert original_logits.shape == logits.shape
   print('''First values of original logits:'''			,       original_logits[0, :3, :3]       )
   print('''First values of HF logits:'''			,       logits[0, :3, :3]       )

   # assert values
   if model_name == "blip2-flan-t5-xl":
      UpperCAmelCase_							:      Tuple						=   torch.tensor(
          [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]			,       device=__lowercase       )
      assert torch.allclose(logits[0, :3, :3]			,       __lowercase			,       atol=1e-4       )
   elif model_name == "blip2-flan-t5-xl-coco":
      UpperCAmelCase_							:      Tuple						=   torch.tensor(
          [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]			,       device=__lowercase       )
   else:
      # cast to same type
      UpperCAmelCase_							:      Optional[int]						=   logits.dtype
      assert torch.allclose(original_logits.to(__lowercase       )			,       __lowercase			,       atol=1e-2       )
   print('''Looks ok!'''       )

   print('''Generating a caption...'''       )
   UpperCAmelCase_							:      Union[str, Any]						=   ''''''
   UpperCAmelCase_							:      Optional[Any]						=   tokenizer(__lowercase			,       return_tensors='''pt'''       ).input_ids.to(__lowercase       )

   UpperCAmelCase_							:      int						=   original_model.generate({'''image''': original_pixel_values}       )
   UpperCAmelCase_							:      Optional[int]						=   hf_model.generate(
       __lowercase			,       __lowercase			,       do_sample=__lowercase			,       num_beams=5			,       max_length=3_0			,       min_length=1			,       top_p=0.9			,       repetition_penalty=1.0			,       length_penalty=1.0			,       temperature=1			,       )
   print('''Original generation:'''			,       __lowercase       )
   UpperCAmelCase_							:      Tuple						=   input_ids.shape[1]
   UpperCAmelCase_							:      str						=   processor.batch_decode(outputs[:, prompt_length:]			,       skip_special_tokens=__lowercase       )
   UpperCAmelCase_							:      Optional[int]						=   [text.strip() for text in output_text]
   print('''HF generation:'''			,       __lowercase       )

   if pytorch_dump_folder_path is not None:
      processor.save_pretrained(__lowercase       )
      hf_model.save_pretrained(__lowercase       )

   if push_to_hub:
      processor.push_to_hub(F'''nielsr/{model_name}'''       )
      hf_model.push_to_hub(F'''nielsr/{model_name}'''       )


if __name__ == "__main__":
       __UpperCamelCase				:  List[Any]      					=							argparse.ArgumentParser()
       __UpperCamelCase				:  Optional[Any]      					=							[
           'blip2-opt-2.7b',
           'blip2-opt-6.7b',
           'blip2-opt-2.7b-coco',
           'blip2-opt-6.7b-coco',
           'blip2-flan-t5-xl',
           'blip2-flan-t5-xl-coco',
           'blip2-flan-t5-xxl',
       ]
       parser.add_argument(
           '--model_name',
           default='blip2-opt-2.7b',
           choices=choices,
           type=str,
           help='Path to hf config.json of model to convert',
       )
       parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
       parser.add_argument(
           '--push_to_hub',
           action='store_true',
           help='Whether to push the model and processor to the hub after converting',
       )

       __UpperCamelCase				:  int      					=							parser.parse_args()

       convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)