


import os
from argparse import ArgumentParser, Namespace

from ..data import SingleSentenceClassificationProcessor as Processor
from ..pipelines import TextClassificationPipeline
from ..utils import is_tf_available, is_torch_available, logging
from . import BaseTransformersCLICommand


if not is_tf_available() and not is_torch_available():
	raise RuntimeError('At least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI training')

# TF training parameters
A_ :				str						 =    False
A_ :				Optional[Any]						 =    False

def 		UpperCamelCase      (lowercase_:				Namespace      )  ->   List[str]:
	return TrainCommand(_UpperCamelCase      )



class 					_a		(__magic_name__       ):


					'''simple docstring'''






					@staticmethod
					def 						__A    (				A__    ):
						A__							:					List[str]       =	parser.add_parser("""train"""					, help="""CLI tool to train a model on a task."""    )

						train_parser.add_argument(
						    """--train_data"""					, type=A__					, required=A__					, help="""path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences."""					, )
						train_parser.add_argument(
						    """--column_label"""					, type=A__					, default=0					, help="""Column of the dataset csv file with example labels."""    )
						train_parser.add_argument(
						    """--column_text"""					, type=A__					, default=1					, help="""Column of the dataset csv file with example texts."""    )
						train_parser.add_argument(
						    """--column_id"""					, type=A__					, default=2					, help="""Column of the dataset csv file with example ids."""    )
						train_parser.add_argument(
						    """--skip_first_row"""					, action="""store_true"""					, help="""Skip the first row of the csv file (headers)."""    )

						train_parser.add_argument("""--validation_data"""					, type=A__					, default=""""""					, help="""path to validation dataset."""    )
						train_parser.add_argument(
						    """--validation_split"""					, type=A__					, default=0.1					, help="""if validation dataset is not provided, fraction of train dataset to use as validation dataset."""					, )

						train_parser.add_argument("""--output"""					, type=A__					, default="""./"""					, help="""path to saved the trained model."""    )

						train_parser.add_argument(
						    """--task"""					, type=A__					, default="""text_classification"""					, help="""Task to train the model on."""    )
						train_parser.add_argument(
						    """--model"""					, type=A__					, default="""bert-base-uncased"""					, help="""Model\'s name or path to stored model."""    )
						train_parser.add_argument("""--train_batch_size"""					, type=A__					, default=32					, help="""Batch size for training."""    )
						train_parser.add_argument("""--valid_batch_size"""					, type=A__					, default=64					, help="""Batch size for validation."""    )
						train_parser.add_argument("""--learning_rate"""					, type=A__					, default=3e-5					, help="""Learning rate."""    )
						train_parser.add_argument("""--adam_epsilon"""					, type=A__					, default=1e-08					, help="""Epsilon for Adam optimizer."""    )
						train_parser.set_defaults(func=A__    )






					def __init__(				self					, A__    ):
						A__							:					str       =	logging.get_logger("""transformers-cli/training"""    )

						A__							:					str       =	"""tf""" if is_tf_available() else """torch"""

						os.makedirs(args.output					, exist_ok=A__    )
						A__							:					Union[str, Any]       =	args.output

						A__							:					Union[str, Any]       =	args.column_label
						A__							:					str       =	args.column_text
						A__							:					List[Any]       =	args.column_id

						self.logger.info(F"""Loading {args.task} pipeline for {args.model}"""    )
						if args.task == "text_classification":
							A__							:					Tuple       =	TextClassificationPipeline.from_pretrained(args.model    )
						elif args.task == "token_classification":
							raise NotImplementedError
						elif args.task == "question_answering":
							raise NotImplementedError

						self.logger.info(F"""Loading dataset from {args.train_data}"""    )
						A__							:					int       =	Processor.create_from_csv(
						    args.train_data					, column_label=args.column_label					, column_text=args.column_text					, column_id=args.column_id					, skip_first_row=args.skip_first_row					, )
						A__							:					Optional[Any]       =	None
						if args.validation_data:
							self.logger.info(F"""Loading validation dataset from {args.validation_data}"""    )
							A__							:					int       =	Processor.create_from_csv(
							    args.validation_data					, column_label=args.column_label					, column_text=args.column_text					, column_id=args.column_id					, skip_first_row=args.skip_first_row					, )

						A__							:					Optional[int]       =	args.validation_split
						A__							:					Optional[int]       =	args.train_batch_size
						A__							:					Optional[Any]       =	args.valid_batch_size
						A__							:					Tuple       =	args.learning_rate
						A__							:					Tuple       =	args.adam_epsilon






					def 						__A    (				self    ):

						if self.framework == "tf":
							return self.run_tf()
						return self.run_torch()






					def 						__A    (				self    ):
						raise NotImplementedError






					def 						__A    (				self    ):
						self.pipeline.fit(
						    self.train_dataset					, validation_data=self.valid_dataset					, validation_split=self.validation_split					, learning_rate=self.learning_rate					, adam_epsilon=self.adam_epsilon					, train_batch_size=self.train_batch_size					, valid_batch_size=self.valid_batch_size					, )

						# Save trained pipeline
						self.pipeline.save_pretrained(self.output    )






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 					_a		(datasets.BeamBasedBuilder       ):


					'''simple docstring'''






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






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






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

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



class 					_a		(datasets.BeamBasedBuilder       ):


					'''simple docstring'''






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






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






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

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



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

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



class 					_a		(__magic_name__       ):


					'''simple docstring'''






					@require_beam
					def 						__A    (				self    ):
						A__							:					Dict       =	len(get_test_dummy_examples()    )
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					int       =	DummyBeamDataset(cache_dir=A__					, beam_runner="""DirectRunner"""    )
							builder.download_and_prepare()
							self.assertTrue(
							    os.path.exists(
							        os.path.join(A__					, builder.name					, """default"""					, """0.0.0"""					, F"""{builder.name}-train.arrow"""    )    )    )
							self.assertDictEqual(builder.info.features					, datasets.Features({"""content""": datasets.Value("""string"""    )}    )    )
							A__							:					int       =	builder.as_dataset()
							self.assertEqual(dset["""train"""].num_rows					, A__    )
							self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples					, A__    )
							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(A__					, builder.name					, """default"""					, """0.0.0"""					, """dataset_info.json"""    )    )    )
							del dset






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

						A__							:					int       =	beam.io.parquetio.WriteToParquet

						A__							:					List[str]       =	len(get_test_dummy_examples()    )
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					str       =	DummyBeamDataset(cache_dir=A__					, beam_runner="""DirectRunner"""    )
							with patch("""apache_beam.io.parquetio.WriteToParquet"""    ) as write_parquet_mock:
								A__							:					Optional[Any]       =	partial(A__					, num_shards=2    )
								builder.download_and_prepare()
							self.assertTrue(
							    os.path.exists(
							        os.path.join(
							            A__					, builder.name					, """default"""					, """0.0.0"""					, F"""{builder.name}-train-00000-of-00002.arrow"""    )    )    )
							self.assertTrue(
							    os.path.exists(
							        os.path.join(
							            A__					, 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"""    )}    )    )
							A__							:					Optional[int]       =	builder.as_dataset()
							self.assertEqual(dset["""train"""].num_rows					, A__    )
							self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples					, A__    )
							# 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(A__					, builder.name					, """default"""					, """0.0.0"""					, """dataset_info.json"""    )    )    )
							del dset






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






					@require_beam
					def 						__A    (				self    ):
						A__							:					List[Any]       =	len(get_test_nested_examples()    )
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					Optional[int]       =	NestedBeamDataset(cache_dir=A__					, beam_runner="""DirectRunner"""    )
							builder.download_and_prepare()
							self.assertTrue(
							    os.path.exists(
							        os.path.join(A__					, 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"""    )}    )}    )    )
							A__							:					Optional[int]       =	builder.as_dataset()
							self.assertEqual(dset["""train"""].num_rows					, A__    )
							self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples					, A__    )
							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(A__					, builder.name					, """default"""					, """0.0.0"""					, """dataset_info.json"""    )    )    )
							del dset



