





"""simple docstring"""






import os
from pathlib import Path
from unittest.mock import patch

import pytest
import zstandard as zstd

from datasets.download.download_config import DownloadConfig
from datasets.utils.file_utils import (
    OfflineModeIsEnabled,
    cached_path,
    fsspec_get,
    fsspec_head,
    ftp_get,
    ftp_head,
    get_from_cache,
    http_get,
    http_head,
)


lowerCAmelCase_			=							'''\
    Text data.
    Second line of data.'''

lowerCAmelCase_			=							'''file'''

@pytest.fixture(scope='session'							)
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Optional[Any]:






       """simple docstring"""




       _UpperCAmelCase     =     tmp_path_factory.mktemp('data'							) / (FILE_PATH + '.zstd')
       _UpperCAmelCase     =     bytes(SCREAMING_SNAKE_CASE,'utf-8'							)
       with zstd.open(SCREAMING_SNAKE_CASE,'wb'							) as f:
              f.write(SCREAMING_SNAKE_CASE							)
       return path

@pytest.fixture
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Tuple:






       """simple docstring"""




       with open(os.path.join(tmpfs.local_root_dir,SCREAMING_SNAKE_CASE							),'w'							) as f:
              f.write(SCREAMING_SNAKE_CASE							)
       return FILE_PATH

@pytest.mark.parametrize('compression_format',['gzip', 'xz', 'zstd']							)
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE							)		->       Any:






       """simple docstring"""




       _UpperCAmelCase     =     {'gzip': gz_file, 'xz': xz_file, 'zstd': zstd_path}
       _UpperCAmelCase     =     input_paths[compression_format]
       _UpperCAmelCase     =     tmp_path / 'cache'
       _UpperCAmelCase     =     DownloadConfig(cache_dir=SCREAMING_SNAKE_CASE,extract_compressed_file=SCREAMING_SNAKE_CASE							)
       _UpperCAmelCase     =     cached_path(SCREAMING_SNAKE_CASE,download_config=SCREAMING_SNAKE_CASE							)
       with open(SCREAMING_SNAKE_CASE							) as f:
              _UpperCAmelCase     =     f.read()
       with open(SCREAMING_SNAKE_CASE							) as f:
              _UpperCAmelCase     =     f.read()
       assert extracted_file_content == expected_file_content

@pytest.mark.parametrize('default_extracted',[True, False]							)
@pytest.mark.parametrize('default_cache_dir',[True, False]							)
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE,SCREAMING_SNAKE_CASE							)		->       int:






       """simple docstring"""




       _UpperCAmelCase     =     'custom_cache'
       _UpperCAmelCase     =     'custom_extracted_dir'
       _UpperCAmelCase     =     tmp_path / 'custom_extracted_path'
       if default_extracted:
              _UpperCAmelCase     =     ('downloads' if default_cache_dir else custom_cache_dir, 'extracted')
       else:
              monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_DIR',SCREAMING_SNAKE_CASE							)
              monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_PATH',str(SCREAMING_SNAKE_CASE							)							)
              _UpperCAmelCase     =     custom_extracted_path.parts[-2:] if default_cache_dir else (custom_cache_dir, custom_extracted_dir)

       _UpperCAmelCase     =     xz_file
       _UpperCAmelCase     =     (
           DownloadConfig(extract_compressed_file=SCREAMING_SNAKE_CASE							)
           if default_cache_dir
           else DownloadConfig(cache_dir=tmp_path / custom_cache_dir,extract_compressed_file=SCREAMING_SNAKE_CASE							)
       )
       _UpperCAmelCase     =     cached_path(SCREAMING_SNAKE_CASE,download_config=SCREAMING_SNAKE_CASE							)
       assert Path(SCREAMING_SNAKE_CASE							).parent.parts[-2:] == expected

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Optional[int]:






       """simple docstring"""




       _UpperCAmelCase     =     str(Path(SCREAMING_SNAKE_CASE							).resolve()							)
       assert cached_path(SCREAMING_SNAKE_CASE							) == text_file
       # relative path
       _UpperCAmelCase     =     str(Path(SCREAMING_SNAKE_CASE							).resolve().relative_to(Path(os.getcwd()							)							)							)
       assert cached_path(SCREAMING_SNAKE_CASE							) == text_file

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Dict:






       """simple docstring"""




       _UpperCAmelCase     =     str(tmp_path.resolve() / '__missing_file__.txt'							)
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              cached_path(SCREAMING_SNAKE_CASE							)
       # relative path
       _UpperCAmelCase     =     './__missing_file__.txt'
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              cached_path(SCREAMING_SNAKE_CASE							)

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       List[str]:






       """simple docstring"""




       _UpperCAmelCase     =     get_from_cache(F"""tmp://{tmpfs_file}"""							)
       with open(SCREAMING_SNAKE_CASE							) as f:
              _UpperCAmelCase     =     f.read()
       assert output_file_content == FILE_CONTENT

@patch('datasets.config.HF_DATASETS_OFFLINE',SCREAMING_SNAKE_CASE							)
def       __lowerCamelCase    (       )		->       List[Any]:






       """simple docstring"""




       with pytest.raises(SCREAMING_SNAKE_CASE							):
              cached_path('https://huggingface.co'							)

@patch('datasets.config.HF_DATASETS_OFFLINE',SCREAMING_SNAKE_CASE							)
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Dict:






       """simple docstring"""




       _UpperCAmelCase     =     tmp_path_factory.mktemp('data'							) / 'file.html'
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              http_get('https://huggingface.co',temp_file=SCREAMING_SNAKE_CASE							)
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              http_head('https://huggingface.co'							)

@patch('datasets.config.HF_DATASETS_OFFLINE',SCREAMING_SNAKE_CASE							)
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       List[Any]:






       """simple docstring"""




       _UpperCAmelCase     =     tmp_path_factory.mktemp('data'							) / 'file.html'
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              ftp_get('ftp://huggingface.co',temp_file=SCREAMING_SNAKE_CASE							)
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              ftp_head('ftp://huggingface.co'							)


@patch('datasets.config.HF_DATASETS_OFFLINE',SCREAMING_SNAKE_CASE							)
def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Dict:






       """simple docstring"""




       _UpperCAmelCase     =     tmp_path_factory.mktemp('data'							) / 'file.html'
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              fsspec_get('s3://huggingface.co',temp_file=SCREAMING_SNAKE_CASE							)
       with pytest.raises(SCREAMING_SNAKE_CASE							):
              fsspec_head('s3://huggingface.co'							)













"""simple docstring"""






import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
    import optax
    from flax.training.common_utils import onehot

    from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
    from transformers.models.ta.modeling_flax_ta import shift_tokens_right



@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 	lowerCAmelCase					(      unittest.TestCase							):






       @slow
       def 							__A (   self     ):
              _UpperCAmelCase     =     FlaxMTaForConditionalGeneration.from_pretrained('google/mt5-small'     )
              _UpperCAmelCase     =     AutoTokenizer.from_pretrained('google/mt5-small'     )

              _UpperCAmelCase     =     tokenizer('Hello there'    ,  return_tensors='np'     ).input_ids
              _UpperCAmelCase     =     tokenizer('Hi I am'    ,  return_tensors='np'     ).input_ids

              _UpperCAmelCase     =     shift_tokens_right(a__    ,  model.config.pad_token_id    ,  model.config.decoder_start_token_id     )

              _UpperCAmelCase     =     model(a__    ,  decoder_input_ids=a__     ).logits
              _UpperCAmelCase     =     optax.softmax_cross_entropy(a__    ,  onehot(a__    ,  logits.shape[-1]     )     ).mean()

              _UpperCAmelCase     =     -(labels.shape[-1] * loss.item())

              _UpperCAmelCase     =     -84.9_127
              self.assertTrue(abs(mtf_score - EXPECTED_SCORE     ) < 1E-4     )






