



import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam



class 	lowerCamelCase      (						datasets.BeamBasedBuilder  ):







					'''simple docstring'''







					def    A__		(			self     ):
							return datasets.DatasetInfo(
							    features=datasets.Features({"content": datasets.Value("string"     )}     )    ,	supervised_keys=lowerCAmelCase    ,	)



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							return [datasets.SplitGenerator(name=datasets.Split.TRAIN    ,	gen_kwargs={"examples": get_test_dummy_examples()}     )]



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							import apache_beam as beam

							return pipeline | "Load Examples" >> beam.Create(lowerCAmelCase     )





class 	lowerCamelCase      (						datasets.BeamBasedBuilder  ):







					'''simple docstring'''







					def    A__		(			self     ):
							return datasets.DatasetInfo(
							    features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"     )}     )}     )    ,	supervised_keys=lowerCAmelCase    ,	)



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							return [
							    datasets.SplitGenerator(name=datasets.Split.TRAIN    ,	gen_kwargs={"examples": get_test_nested_examples()}     )
							]



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							import apache_beam as beam

							return pipeline | "Load Examples" >> beam.Create(lowerCAmelCase     )







def snake_case__  (			)		->							Tuple:
		return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"] )]







def snake_case__  (			)		->							Any:
		return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["foo", "bar", "foobar"] )]



class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					@require_beam
					def    A__		(			self     ):
							UpperCAmelCase_						=				len(get_test_dummy_examples()     )
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
									UpperCAmelCase_						=				DummyBeamDataset(cache_dir=lowerCAmelCase    ,	beam_runner="DirectRunner"     )
									builder.download_and_prepare()
									self.assertTrue(
									    os.path.exists(
									        os.path.join(lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	f'''{builder.name}-train.arrow'''     )     )     )
									self.assertDictEqual(builder.info.features    ,	datasets.Features({"content": datasets.Value("string"     )}     )     )
									UpperCAmelCase_						=				builder.as_dataset()
									self.assertEqual(dset["train"].num_rows    ,	lowerCAmelCase     )
									self.assertEqual(dset["train"].info.splits["train"].num_examples    ,	lowerCAmelCase     )
									self.assertDictEqual(dset["train"][0]    ,	get_test_dummy_examples()[0][1]     )
									self.assertDictEqual(
									    dset["train"][expected_num_examples - 1]    ,	get_test_dummy_examples()[expected_num_examples - 1][1]     )
									self.assertTrue(
									    os.path.exists(os.path.join(lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	"dataset_info.json"     )     )     )
									del dset



					@require_beam
					def    A__		(			self     ):
							import apache_beam as beam

							UpperCAmelCase_						=				beam.io.parquetio.WriteToParquet

							UpperCAmelCase_						=				len(get_test_dummy_examples()     )
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
									UpperCAmelCase_						=				DummyBeamDataset(cache_dir=lowerCAmelCase    ,	beam_runner="DirectRunner"     )
									with patch("apache_beam.io.parquetio.WriteToParquet"     ) as write_parquet_mock:
											UpperCAmelCase_						=				partial(lowerCAmelCase    ,	num_shards=2     )
											builder.download_and_prepare()
									self.assertTrue(
									    os.path.exists(
									        os.path.join(
									            lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	f'''{builder.name}-train-00000-of-00002.arrow'''     )     )     )
									self.assertTrue(
									    os.path.exists(
									        os.path.join(
									            lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	f'''{builder.name}-train-00000-of-00002.arrow'''     )     )     )
									self.assertDictEqual(builder.info.features    ,	datasets.Features({"content": datasets.Value("string"     )}     )     )
									UpperCAmelCase_						=				builder.as_dataset()
									self.assertEqual(dset["train"].num_rows    ,	lowerCAmelCase     )
									self.assertEqual(dset["train"].info.splits["train"].num_examples    ,	lowerCAmelCase     )
									# Order is not preserved when sharding, so we just check that all the elements are there
									self.assertListEqual(sorted(dset["train"]["content"]     )    ,	sorted(["foo", "bar", "foobar"]     )     )
									self.assertTrue(
									    os.path.exists(os.path.join(lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	"dataset_info.json"     )     )     )
									del dset



					@require_beam
					def    A__		(			self     ):
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
									UpperCAmelCase_						=				DummyBeamDataset(cache_dir=lowerCAmelCase     )
									self.assertRaises(datasets.builder.MissingBeamOptions    ,	builder.download_and_prepare     )



					@require_beam
					def    A__		(			self     ):
							UpperCAmelCase_						=				len(get_test_nested_examples()     )
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
									UpperCAmelCase_						=				NestedBeamDataset(cache_dir=lowerCAmelCase    ,	beam_runner="DirectRunner"     )
									builder.download_and_prepare()
									self.assertTrue(
									    os.path.exists(
									        os.path.join(lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	f'''{builder.name}-train.arrow'''     )     )     )
									self.assertDictEqual(
									    builder.info.features    ,	datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"     )}     )}     )     )
									UpperCAmelCase_						=				builder.as_dataset()
									self.assertEqual(dset["train"].num_rows    ,	lowerCAmelCase     )
									self.assertEqual(dset["train"].info.splits["train"].num_examples    ,	lowerCAmelCase     )
									self.assertDictEqual(dset["train"][0]    ,	get_test_nested_examples()[0][1]     )
									self.assertDictEqual(
									    dset["train"][expected_num_examples - 1]    ,	get_test_nested_examples()[expected_num_examples - 1][1]     )
									self.assertTrue(
									    os.path.exists(os.path.join(lowerCAmelCase    ,	builder.name    ,	"default"    ,	"0.0.0"    ,	"dataset_info.json"     )     )     )
									del dset






import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
			from transformers import AutoModelForSeqaSeqLM, AutoTokenizer



@require_torch
@require_sentencepiece
@require_tokenizers
class 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







					@slow
					def    A__		(			self     ):
							UpperCAmelCase_						=				AutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"    ,	return_dict=lowerCAmelCase     ).to(lowerCAmelCase     )
							UpperCAmelCase_						=				AutoTokenizer.from_pretrained("google/mt5-small"     )

							UpperCAmelCase_						=				tokenizer("Hello there"    ,	return_tensors="pt"     ).input_ids
							UpperCAmelCase_						=				tokenizer("Hi I am"    ,	return_tensors="pt"     ).input_ids

							UpperCAmelCase_						=				model(input_ids.to(lowerCAmelCase     )    ,	labels=labels.to(lowerCAmelCase     )     ).loss
							UpperCAmelCase_						=				-(labels.shape[-1] * loss.item())

							UpperCAmelCase_						=				-84.9127
							self.assertTrue(abs(mtf_score - EXPECTED_SCORE     ) < 1e-4     )

