


import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
UpperCamelCase								=	logging.getLogger()

UpperCamelCase								=	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)






class   _lowerCamelCase (					UpperCamelCase			):




       """simple docstring"""



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE		)->Optional[Any]:




         '''simple docstring'''





         os.makedirs(_SCREAMING_SNAKE_CASE				,       exist_ok=_SCREAMING_SNAKE_CASE		)
         A_      :				Tuple		     =					{'''source''': '''What is love ?''', '''target''': '''life'''}
         A_      :				List[str]		     =					{'''train''': 12, '''val''': 2, '''test''': 2}
         for split in ["train", "test", "val"]:
           for field in ["source", "target"]:
             A_      :				Any		     =					'''\n'''.join([contents[field]] * n_lines[split]		)
             with open(os.path.join(_SCREAMING_SNAKE_CASE				,       F'''{split}.{field}'''		)				,       '''w'''		) as f:
               f.write(_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE = "pytorch"		)->List[str]:




         '''simple docstring'''





         A_      :				int		     =					self.get_auto_remove_tmp_dir()
         A_      :				str		     =					os.path.join(_SCREAMING_SNAKE_CASE				,       '''output'''		)
         A_      :				Tuple		     =					os.path.join(_SCREAMING_SNAKE_CASE				,       '''data'''		)
         self._create_dummy_data(data_dir=_SCREAMING_SNAKE_CASE		)

         A_      :				Union[str, Any]		     =					F'''
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            '''.split()

         if gpus > 0:
           testargs.append(F'''--gpus={gpus}'''		)
           if is_apex_available():
             testargs.append('''--fp16'''		)
         else:
           testargs.append('''--gpus=0'''		)
           testargs.append('''--distributed_backend=ddp_cpu'''		)
           testargs.append('''--num_processes=2'''		)

         A_      :				Optional[int]		     =					[sys.executable, str(Path(finetune_rag.__file__		).resolve()		)] + testargs
         execute_subprocess_async(_SCREAMING_SNAKE_CASE				,       env=self.get_env()		)

         A_      :				Optional[Any]		     =					os.path.join(_SCREAMING_SNAKE_CASE				,       '''metrics.json'''		)
         with open(_SCREAMING_SNAKE_CASE		) as f:
           A_      :				int		     =					json.load(_SCREAMING_SNAKE_CASE		)
         return result



       @require_torch_gpu
       def        _snake_case       (      self		)->Any:




         '''simple docstring'''





         A_      :				str		     =					self._run_finetune(gpus=1		)
         self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']				,       0.2		)



       @require_torch_multi_gpu
       def        _snake_case       (      self		)->List[Any]:




         '''simple docstring'''





         A_      :				Tuple		     =					self._run_finetune(gpus=2		)
         self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']				,       0.2		)



       @require_torch_gpu
       @require_ray
       def        _snake_case       (      self		)->Union[str, Any]:




         '''simple docstring'''





         A_      :				Tuple		     =					self._run_finetune(gpus=1				,       distributed_retriever='''ray'''		)
         self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']				,       0.2		)







       @require_torch_multi_gpu
       @require_ray
       def        _snake_case       (      self		)->Dict:




         '''simple docstring'''





         A_      :				List[Any]		     =					self._run_finetune(gpus=1				,       distributed_retriever='''ray'''		)
         self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']				,       0.2		)








import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  A_      :				Any		     =					np.inf

  def set_batch_size(SCREAMING_SNAKE_CASE       ) -> None:
    nonlocal batch_size
    if isinstance(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
      A_      :				Tuple		     =					min(SCREAMING_SNAKE_CASE ,			config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS       )
    elif isinstance(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
      A_      :				str		     =					min(SCREAMING_SNAKE_CASE ,			config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS       )
    elif isinstance(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ) and feature.dtype == "binary":
      A_      :				Union[str, Any]		     =					min(SCREAMING_SNAKE_CASE ,			config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS       )

  _visit(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       )

  return None if batch_size is np.inf else batch_size






class   _lowerCamelCase (					UpperCamelCase			):




       """simple docstring"""



       def __init__(      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE = None				,       _SCREAMING_SNAKE_CASE = None				,       _SCREAMING_SNAKE_CASE = None				,       _SCREAMING_SNAKE_CASE = False				,       _SCREAMING_SNAKE_CASE = False				,       _SCREAMING_SNAKE_CASE = None				,       **_SCREAMING_SNAKE_CASE				,       )->str:




         '''simple docstring'''





         super().__init__(
             _SCREAMING_SNAKE_CASE				,       split=_SCREAMING_SNAKE_CASE				,       features=_SCREAMING_SNAKE_CASE				,       cache_dir=_SCREAMING_SNAKE_CASE				,       keep_in_memory=_SCREAMING_SNAKE_CASE				,       streaming=_SCREAMING_SNAKE_CASE				,       num_proc=_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE				,       )
         A_      :				str		     =					path_or_paths if isinstance(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		) else {self.split: path_or_paths}
         A_      :				Optional[int]		     =					_PACKAGED_DATASETS_MODULES['''parquet'''][1]
         A_      :				Union[str, Any]		     =					Parquet(
             cache_dir=_SCREAMING_SNAKE_CASE				,       data_files=_SCREAMING_SNAKE_CASE				,       features=_SCREAMING_SNAKE_CASE				,       hash=_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE				,       )







       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





         if self.streaming:
           A_      :				List[str]		     =					self.builder.as_streaming_dataset(split=self.split		)
         # Build regular (map-style) dataset
         else:
           A_      :				List[str]		     =					None
           A_      :				List[str]		     =					None
           A_      :				List[Any]		     =					None
           A_      :				Dict		     =					None

           self.builder.download_and_prepare(
               download_config=_SCREAMING_SNAKE_CASE				,       download_mode=_SCREAMING_SNAKE_CASE				,       verification_mode=_SCREAMING_SNAKE_CASE				,       base_path=_SCREAMING_SNAKE_CASE				,       num_proc=self.num_proc				,       )
           A_      :				Union[str, Any]		     =					self.builder.as_dataset(
               split=self.split				,       verification_mode=_SCREAMING_SNAKE_CASE				,       in_memory=self.keep_in_memory		)
         return dataset






class   _lowerCamelCase :




       """simple docstring"""



       def __init__(      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE = None				,       **_SCREAMING_SNAKE_CASE				,       )->List[Any]:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					dataset
         A_      :				Union[str, Any]		     =					path_or_buf
         A_      :				Any		     =					batch_size or get_writer_batch_size(dataset.features		)
         A_      :				Optional[int]		     =					parquet_writer_kwargs



       def        _snake_case       (      self		)->int:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

         if isinstance(self.path_or_buf				,       (str, bytes, os.PathLike)		):
           with open(self.path_or_buf				,       '''wb+'''		) as buffer:
             A_      :				str		     =					self._write(file_obj=_SCREAMING_SNAKE_CASE				,       batch_size=_SCREAMING_SNAKE_CASE				,       **self.parquet_writer_kwargs		)
         else:
           A_      :				Tuple		     =					self._write(file_obj=self.path_or_buf				,       batch_size=_SCREAMING_SNAKE_CASE				,       **self.parquet_writer_kwargs		)
         return written







       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)->int:




         '''simple docstring'''





         A_      :				List[Any]		     =					0
         A_      :				int		     =					parquet_writer_kwargs.pop('''path_or_buf'''				,       _SCREAMING_SNAKE_CASE		)
         A_      :				List[Any]		     =					self.dataset.features.arrow_schema

         A_      :				List[str]		     =					pq.ParquetWriter(_SCREAMING_SNAKE_CASE				,       schema=_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)

         for offset in logging.tqdm(
             range(0				,       len(self.dataset		)				,       _SCREAMING_SNAKE_CASE		)				,       unit='''ba'''				,       disable=not logging.is_progress_bar_enabled()				,       desc='''Creating parquet from Arrow format'''				,       ):
           A_      :				List[Any]		     =					query_table(
               table=self.dataset._data				,       key=slice(_SCREAMING_SNAKE_CASE				,       offset + batch_size		)				,       indices=self.dataset._indices if self.dataset._indices is not None else None				,       )
           writer.write_table(_SCREAMING_SNAKE_CASE		)
           written += batch.nbytes
         writer.close()
         return written




