

import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


lowerCAmelCase							:						List[Any]     			=		logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

lowerCAmelCase							:						int     			=		{
    """base""": AutoModel,
    """sequence-classification""": AutoModelForSequenceClassification,
    """question-answering""": AutoModelForQuestionAnswering,
    """pretraining""": AutoModelForPreTraining,
    """token-classification""": AutoModelForTokenClassification,
    """language-modeling""": AutoModelWithLMHead,
    """summarization""": AutoModelForSeqaSeqLM,
    """translation""": AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
lowerCAmelCase							:						Any     			=		{
    """linear""": get_linear_schedule_with_warmup,
    """cosine""": get_cosine_schedule_with_warmup,
    """cosine_w_restarts""": get_cosine_with_hard_restarts_schedule_with_warmup,
    """polynomial""": get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
lowerCAmelCase							:						Dict     			=		sorted(arg_to_scheduler.keys())
lowerCAmelCase							:						List[Any]     			=		"""{""" + """, """.join(arg_to_scheduler_choices) + """}"""



class 						__lowercase       (				pl.LightningModule			):




   """simple docstring"""







   def __init__(				self						:  int     ,		lowerCAmelCase__						:  Dict     ,		lowerCAmelCase__						:  Union[str, Any]=None     ,		lowerCAmelCase__						:  Tuple="base"     ,		lowerCAmelCase__						:  Dict=None     ,		lowerCAmelCase__						:  List[str]=None     ,		lowerCAmelCase__						:  Union[str, Any]=None     ,		**lowerCAmelCase__						:  Any     ,		):
         super().__init__()
         # TODO: move to self.save_hyperparameters()
         # self.save_hyperparameters()
         # can also expand arguments into trainer signature for easier reading

         self.save_hyperparameters(__lowerCamelCase)
         SCREAMING_SNAKE_CASE_:				List[str]							= 0
         SCREAMING_SNAKE_CASE_:				Dict							= Path(self.hparams.output_dir)
         SCREAMING_SNAKE_CASE_:				Optional[Any]							= self.hparams.cache_dir if self.hparams.cache_dir else None
         if config is None:
               SCREAMING_SNAKE_CASE_:				str							= AutoConfig.from_pretrained(
                   self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path     ,		**({"num_labels": num_labels} if num_labels is not None else {})     ,		cache_dir=__lowerCamelCase     ,		**__lowerCamelCase     ,		)
         else:
               SCREAMING_SNAKE_CASE_:				PretrainedConfig							= config

         SCREAMING_SNAKE_CASE_:				int							= ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
         for p in extra_model_params:
               if getattr(self.hparams     ,		__lowerCamelCase     ,		__lowerCamelCase):
                     assert hasattr(self.config     ,		__lowerCamelCase), F"model config doesn't have a `{p}` attribute"
                     setattr(self.config     ,		__lowerCamelCase     ,		getattr(self.hparams     ,		__lowerCamelCase))

         if tokenizer is None:
               SCREAMING_SNAKE_CASE_:				Union[str, Any]							= AutoTokenizer.from_pretrained(
                   self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path     ,		cache_dir=__lowerCamelCase     ,		)
         else:
               SCREAMING_SNAKE_CASE_:				PreTrainedTokenizer							= tokenizer
         SCREAMING_SNAKE_CASE_:				Any							= MODEL_MODES[mode]
         if model is None:
               SCREAMING_SNAKE_CASE_:				Optional[Any]							= self.model_type.from_pretrained(
                   self.hparams.model_name_or_path     ,		from_tf=bool(".ckpt" in self.hparams.model_name_or_path)     ,		config=self.config     ,		cache_dir=__lowerCamelCase     ,		)
         else:
               SCREAMING_SNAKE_CASE_:				Optional[Any]							= model



   def  _SCREAMING_SNAKE_CASE    (				self						:  Tuple     ,		*lowerCAmelCase__						:  List[str]     ,		**lowerCAmelCase__						:  Tuple):
         SCREAMING_SNAKE_CASE_:				int							= self.model_type.from_pretrained(*__lowerCamelCase     ,		**__lowerCamelCase)



   def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[int]):
         SCREAMING_SNAKE_CASE_:				Dict							= arg_to_scheduler[self.hparams.lr_scheduler]
         SCREAMING_SNAKE_CASE_:				int							= get_schedule_func(
             self.opt     ,		num_warmup_steps=self.hparams.warmup_steps     ,		num_training_steps=self.total_steps())
         SCREAMING_SNAKE_CASE_:				Union[str, Any]							= {"scheduler": scheduler, "interval": "step", "frequency": 1}
         return scheduler



   def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[Any]):
         SCREAMING_SNAKE_CASE_:				Dict							= self.model
         SCREAMING_SNAKE_CASE_:				Optional[Any]							= ["bias", "LayerNorm.weight"]
         SCREAMING_SNAKE_CASE_:				Tuple							= [
             {
                 "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                 "weight_decay": self.hparams.weight_decay,
             },
             {
                 "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
                 "weight_decay": 0.0,
             },
         ]
         if self.hparams.adafactor:
               SCREAMING_SNAKE_CASE_:				Optional[int]							= Adafactor(
                   __lowerCamelCase     ,		lr=self.hparams.learning_rate     ,		scale_parameter=__lowerCamelCase     ,		relative_step=__lowerCamelCase)

         else:
               SCREAMING_SNAKE_CASE_:				str							= AdamW(
                   __lowerCamelCase     ,		lr=self.hparams.learning_rate     ,		eps=self.hparams.adam_epsilon)
         SCREAMING_SNAKE_CASE_:				Tuple							= optimizer

         SCREAMING_SNAKE_CASE_:				List[str]							= self.get_lr_scheduler()

         return [optimizer], [scheduler]



   def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]     ,		lowerCAmelCase__						:  List[str]     ,		lowerCAmelCase__						:  Union[str, Any]):
         return self.validation_step(__lowerCamelCase     ,		__lowerCamelCase)



   def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[int]     ,		lowerCAmelCase__						:  Union[str, Any]):
         return self.validation_end(__lowerCamelCase)



   def  _SCREAMING_SNAKE_CASE    (				self						:  str):
         SCREAMING_SNAKE_CASE_:				List[Any]							= max(1     ,		self.hparams.gpus)  # TODO: consider num_tpu_cores
         SCREAMING_SNAKE_CASE_:				Optional[int]							= self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
         return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



   def  _SCREAMING_SNAKE_CASE    (				self						:  List[str]     ,		lowerCAmelCase__						:  str):
         if mode == "test":
               SCREAMING_SNAKE_CASE_:				Tuple							= len(self.test_dataloader().dataset)
         else:
               SCREAMING_SNAKE_CASE_:				str							= self.get_dataloader("train"     ,		self.hparams.train_batch_size     ,		shuffle=__lowerCamelCase)
               SCREAMING_SNAKE_CASE_:				Any							= len(self.train_dataloader().dataset)



   def  _SCREAMING_SNAKE_CASE    (				self						:  Tuple     ,		lowerCAmelCase__						:  List[Any]     ,		lowerCAmelCase__						:  Optional[int]     ,		lowerCAmelCase__						:  str = False):
         raise NotImplementedError("You must implement this for your task")



   def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]):
         return self.train_loader



   def  _SCREAMING_SNAKE_CASE    (				self						:  List[str]):
         return self.get_dataloader("dev"     ,		self.hparams.eval_batch_size     ,		shuffle=__lowerCamelCase)



   def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[Any]):
         return self.get_dataloader("test"     ,		self.hparams.eval_batch_size     ,		shuffle=__lowerCamelCase)



   def  _SCREAMING_SNAKE_CASE    (				self						:  Any     ,		lowerCAmelCase__						:  str):
         return os.path.join(
             self.hparams.data_dir     ,		"cached_{}_{}_{}".format(
                 __lowerCamelCase     ,		list(filter(__lowerCamelCase     ,		self.hparams.model_name_or_path.split("/"))).pop()     ,		str(self.hparams.max_seq_length)     ,		)     ,		)



   @pl.utilities.rank_zero_only
   def  _SCREAMING_SNAKE_CASE    (				self						:  Union[str, Any]     ,		lowerCAmelCase__						:  Dict):
         SCREAMING_SNAKE_CASE_:				Any							= self.output_dir.joinpath("best_tfmr")
         SCREAMING_SNAKE_CASE_:				str							= self.step_count
         self.model.save_pretrained(__lowerCamelCase)
         self.tokenizer.save_pretrained(__lowerCamelCase)




   @staticmethod
   def  _SCREAMING_SNAKE_CASE    (				lowerCAmelCase__						:  Dict     ,		lowerCAmelCase__						:  Any):
         parser.add_argument(
             "--model_name_or_path"     ,		default=__lowerCamelCase     ,		type=__lowerCamelCase     ,		required=__lowerCamelCase     ,		help="Path to pretrained model or model identifier from huggingface.co/models"     ,		)
         parser.add_argument(
             "--config_name"     ,		default=""     ,		type=__lowerCamelCase     ,		help="Pretrained config name or path if not the same as model_name")
         parser.add_argument(
             "--tokenizer_name"     ,		default=__lowerCamelCase     ,		type=__lowerCamelCase     ,		help="Pretrained tokenizer name or path if not the same as model_name"     ,		)
         parser.add_argument(
             "--cache_dir"     ,		default=""     ,		type=__lowerCamelCase     ,		help="Where do you want to store the pre-trained models downloaded from huggingface.co"     ,		)
         parser.add_argument(
             "--encoder_layerdrop"     ,		type=__lowerCamelCase     ,		help="Encoder layer dropout probability (Optional). Goes into model.config"     ,		)
         parser.add_argument(
             "--decoder_layerdrop"     ,		type=__lowerCamelCase     ,		help="Decoder layer dropout probability (Optional). Goes into model.config"     ,		)
         parser.add_argument(
             "--dropout"     ,		type=__lowerCamelCase     ,		help="Dropout probability (Optional). Goes into model.config"     ,		)
         parser.add_argument(
             "--attention_dropout"     ,		type=__lowerCamelCase     ,		help="Attention dropout probability (Optional). Goes into model.config"     ,		)
         parser.add_argument("--learning_rate"     ,		default=5E-5     ,		type=__lowerCamelCase     ,		help="The initial learning rate for Adam.")
         parser.add_argument(
             "--lr_scheduler"     ,		default="linear"     ,		choices=__lowerCamelCase     ,		metavar=__lowerCamelCase     ,		type=__lowerCamelCase     ,		help="Learning rate scheduler"     ,		)
         parser.add_argument("--weight_decay"     ,		default=0.0     ,		type=__lowerCamelCase     ,		help="Weight decay if we apply some.")
         parser.add_argument("--adam_epsilon"     ,		default=1E-8     ,		type=__lowerCamelCase     ,		help="Epsilon for Adam optimizer.")
         parser.add_argument("--warmup_steps"     ,		default=0     ,		type=__lowerCamelCase     ,		help="Linear warmup over warmup_steps.")
         parser.add_argument("--num_workers"     ,		default=4     ,		type=__lowerCamelCase     ,		help="kwarg passed to DataLoader")
         parser.add_argument("--num_train_epochs"     ,		dest="max_epochs"     ,		default=3     ,		type=__lowerCamelCase)
         parser.add_argument("--train_batch_size"     ,		default=32     ,		type=__lowerCamelCase)
         parser.add_argument("--eval_batch_size"     ,		default=32     ,		type=__lowerCamelCase)
         parser.add_argument("--adafactor"     ,		action="store_true")







class 						__lowercase       (				pl.Callback			):




   """simple docstring"""







   def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]     ,		lowerCAmelCase__						:  int     ,		lowerCAmelCase__						:  List[Any]):
         SCREAMING_SNAKE_CASE_:				Optional[Any]							= trainer.lr_schedulers[0]["scheduler"]
         SCREAMING_SNAKE_CASE_:				Tuple							= {F"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr())}
         pl_module.logger.log_metrics(__lowerCamelCase)



   def  _SCREAMING_SNAKE_CASE    (				self						:  Any     ,		lowerCAmelCase__						:  List[Any]     ,		lowerCAmelCase__						:  Optional[int]):
         rank_zero_info("***** Validation results *****")
         SCREAMING_SNAKE_CASE_:				Optional[int]							= trainer.callback_metrics
         # Log results
         for key in sorted(__lowerCamelCase):
               if key not in ["log", "progress_bar"]:
                     rank_zero_info("{} = {}\n".format(__lowerCamelCase     ,		str(metrics[key])))




   def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[int]     ,		lowerCAmelCase__						:  List[str]     ,		lowerCAmelCase__						:  Tuple):
         rank_zero_info("***** Test results *****")
         SCREAMING_SNAKE_CASE_:				Tuple							= trainer.callback_metrics
         # Log and save results to file
         SCREAMING_SNAKE_CASE_:				str							= os.path.join(pl_module.hparams.output_dir     ,		"test_results.txt")
         with open(__lowerCamelCase     ,		"w") as writer:
               for key in sorted(__lowerCamelCase):
                     if key not in ["log", "progress_bar"]:
                           rank_zero_info("{} = {}\n".format(__lowerCamelCase     ,		str(metrics[key])))
                           writer.write("{} = {}\n".format(__lowerCamelCase     ,		str(metrics[key])))







def     A_   (					_UpperCAmelCase       ,			_UpperCAmelCase					):
      #  To allow all pl args uncomment the following line
      #  parser = pl.Trainer.add_argparse_args(parser)
      parser.add_argument(
          "--output_dir"       ,			default=UpperCamelCase__       ,			type=UpperCamelCase__       ,			required=UpperCamelCase__       ,			help="The output directory where the model predictions and checkpoints will be written."       ,			)
      parser.add_argument(
          "--fp16"       ,			action="store_true"       ,			help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"       ,			)

      parser.add_argument(
          "--fp16_opt_level"       ,			type=UpperCamelCase__       ,			default="O2"       ,			help=(
              "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
              "See details at https://nvidia.github.io/apex/amp.html"
          )       ,			)
      parser.add_argument("--n_tpu_cores"       ,			dest="tpu_cores"       ,			type=UpperCamelCase__					)
      parser.add_argument("--max_grad_norm"       ,			dest="gradient_clip_val"       ,			default=1.0       ,			type=UpperCamelCase__       ,			help="Max gradient norm"					)
      parser.add_argument("--do_train"       ,			action="store_true"       ,			help="Whether to run training."					)
      parser.add_argument("--do_predict"       ,			action="store_true"       ,			help="Whether to run predictions on the test set."					)
      parser.add_argument(
          "--gradient_accumulation_steps"       ,			dest="accumulate_grad_batches"       ,			type=UpperCamelCase__       ,			default=1       ,			help="Number of updates steps to accumulate before performing a backward/update pass."       ,			)
      parser.add_argument("--seed"       ,			type=UpperCamelCase__       ,			default=42       ,			help="random seed for initialization"					)
      parser.add_argument(
          "--data_dir"       ,			default=UpperCamelCase__       ,			type=UpperCamelCase__       ,			required=UpperCamelCase__       ,			help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."       ,			)







def     A_   (					_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=True       ,			_UpperCAmelCase=[]       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=None       ,			**_UpperCAmelCase       ,			):
      pl.seed_everything(args.seed					)

      # init model
      SCREAMING_SNAKE_CASE_:				Dict							= Path(model.hparams.output_dir					)
      odir.mkdir(exist_ok=UpperCamelCase__					)

      # add custom checkpoints
      if checkpoint_callback is None:
            SCREAMING_SNAKE_CASE_:				Optional[int]							= pl.callbacks.ModelCheckpoint(
                filepath=args.output_dir       ,			prefix="checkpoint"       ,			monitor="val_loss"       ,			mode="min"       ,			save_top_k=1					)
      if early_stopping_callback:
            extra_callbacks.append(UpperCamelCase__					)
      if logging_callback is None:
            SCREAMING_SNAKE_CASE_:				List[Any]							= LoggingCallback()

      SCREAMING_SNAKE_CASE_:				str							= {}

      # TODO: remove with PyTorch 1.6 since pl uses native amp
      if args.fpaa:
            SCREAMING_SNAKE_CASE_:				str							= 16
            SCREAMING_SNAKE_CASE_:				List[Any]							= args.fpaa_opt_level

      if args.gpus > 1:
            SCREAMING_SNAKE_CASE_:				int							= "ddp"

      SCREAMING_SNAKE_CASE_:				Optional[int]							= args.accumulate_grad_batches
      SCREAMING_SNAKE_CASE_:				Optional[Any]							= extra_train_kwargs.get("accelerator"       ,			UpperCamelCase__					)
      SCREAMING_SNAKE_CASE_:				List[str]							= extra_train_kwargs.get("profiler"       ,			UpperCamelCase__					)

      SCREAMING_SNAKE_CASE_:				List[Any]							= pl.Trainer.from_argparse_args(
          UpperCamelCase__       ,			weights_summary=UpperCamelCase__       ,			callbacks=[logging_callback] + extra_callbacks       ,			logger=UpperCamelCase__       ,			checkpoint_callback=UpperCamelCase__       ,			**UpperCamelCase__       ,			)

      if args.do_train:
            trainer.fit(UpperCamelCase__					)

      return trainer









import itertools
from dataclasses import dataclass
from io import StringIO
from typing import Optional

import pyarrow as pa

import datasets
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


lowerCAmelCase							:						Tuple     			=		datasets.utils.logging.get_logger(__name__)



@dataclass
class 						__lowercase       (				datasets.BuilderConfig			):




			"""simple docstring"""







			_UpperCAmelCase     :							Optional[datasets.Features]    							=  None
			_UpperCAmelCase     :							str    							=  "utf-8"
			_UpperCAmelCase     :							Optional[str]    							=  None
			_UpperCAmelCase     :							int    							=  10 << 20  # 10MB
			_UpperCAmelCase     :							bool    							=  False
			_UpperCAmelCase     :							str    							=  "line"







class 						__lowercase       (				datasets.ArrowBasedBuilder			):




			"""simple docstring"""







			_UpperCAmelCase     :							List[Any]    							=  TextConfig



			def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[Any]):
									return datasets.DatasetInfo(features=self.config.features)



			def  _SCREAMING_SNAKE_CASE    (				self						:  List[Any]     ,		lowerCAmelCase__						:  Tuple):
									if not self.config.data_files:
															raise ValueError(F"At least one data file must be specified, but got data_files={self.config.data_files}")
									SCREAMING_SNAKE_CASE_:				Dict							= dl_manager.download_and_extract(self.config.data_files)
									if isinstance(lowerCAmelCase__     ,		(str, list, tuple)):
															SCREAMING_SNAKE_CASE_:				Dict							= data_files
															if isinstance(lowerCAmelCase__     ,		lowerCAmelCase__):
																					SCREAMING_SNAKE_CASE_:				Dict							= [files]
															SCREAMING_SNAKE_CASE_:				List[str]							= [dl_manager.iter_files(lowerCAmelCase__) for file in files]
															return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,		gen_kwargs={"files": files})]
									SCREAMING_SNAKE_CASE_:				Any							= []
									for split_name, files in data_files.items():
															if isinstance(lowerCAmelCase__     ,		lowerCAmelCase__):
																					SCREAMING_SNAKE_CASE_:				List[Any]							= [files]
															SCREAMING_SNAKE_CASE_:				Union[str, Any]							= [dl_manager.iter_files(lowerCAmelCase__) for file in files]
															splits.append(datasets.SplitGenerator(name=lowerCAmelCase__     ,		gen_kwargs={"files": files}))
									return splits



			def  _SCREAMING_SNAKE_CASE    (				self						:  Optional[int]     ,		lowerCAmelCase__						:  pa.Table):
									if self.config.features is not None:
															SCREAMING_SNAKE_CASE_:				Optional[int]							= self.config.features.arrow_schema
															if all(not require_storage_cast(lowerCAmelCase__) for feature in self.config.features.values()):
																					# cheaper cast
																					SCREAMING_SNAKE_CASE_:				List[str]							= pa_table.cast(lowerCAmelCase__)
															else:
																					# more expensive cast; allows str <-> int/float or str to Audio for example
																					SCREAMING_SNAKE_CASE_:				List[Any]							= table_cast(lowerCAmelCase__     ,		lowerCAmelCase__)
															return pa_table
									else:
															return pa_table.cast(pa.schema({"text": pa.string()}))




			def  _SCREAMING_SNAKE_CASE    (				self						:  int     ,		lowerCAmelCase__						:  Dict):
									SCREAMING_SNAKE_CASE_:				str							= list(self.config.features) if self.config.features is not None else ["text"]
									for file_idx, file in enumerate(itertools.chain.from_iterable(lowerCAmelCase__)):
															# open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
															with open(lowerCAmelCase__     ,		encoding=self.config.encoding     ,		errors=self.config.errors) as f:
																					if self.config.sample_by == "line":
																											SCREAMING_SNAKE_CASE_:				Any							= 0
																											while True:
																																	SCREAMING_SNAKE_CASE_:				str							= f.read(self.config.chunksize)
																																	if not batch:
																																							break
																																	batch += f.readline()  # finish current line
																																	# StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
																																	SCREAMING_SNAKE_CASE_:				Optional[Any]							= StringIO(lowerCAmelCase__).readlines()
																																	if not self.config.keep_linebreaks:
																																							SCREAMING_SNAKE_CASE_:				Optional[Any]							= [line.rstrip("\n") for line in batch]
																																	SCREAMING_SNAKE_CASE_:				Any							= pa.Table.from_arrays([pa.array(lowerCAmelCase__)]     ,		names=lowerCAmelCase__)
																																	# Uncomment for debugging (will print the Arrow table size and elements)
																																	# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																																	# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																	yield (file_idx, batch_idx), self._cast_table(lowerCAmelCase__)
																																	batch_idx += 1
																					elif self.config.sample_by == "paragraph":
																											SCREAMING_SNAKE_CASE_:				List[str]							= 0
																											SCREAMING_SNAKE_CASE_:				Tuple							= ""
																											while True:
																																	SCREAMING_SNAKE_CASE_:				Dict							= f.read(self.config.chunksize)
																																	if not new_batch:
																																							break
																																	batch += new_batch
																																	batch += f.readline()  # finish current line
																																	SCREAMING_SNAKE_CASE_:				List[str]							= batch.split("\n\n")
																																	SCREAMING_SNAKE_CASE_:				Optional[int]							= pa.Table.from_arrays(
																																	    [pa.array([example for example in batch[:-1] if example])]     ,		names=lowerCAmelCase__)
																																	# Uncomment for debugging (will print the Arrow table size and elements)
																																	# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																																	# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																	yield (file_idx, batch_idx), self._cast_table(lowerCAmelCase__)
																																	batch_idx += 1
																																	SCREAMING_SNAKE_CASE_:				Optional[Any]							= batch[-1]
																											if batch:
																																	SCREAMING_SNAKE_CASE_:				int							= pa.Table.from_arrays([pa.array([batch])]     ,		names=lowerCAmelCase__)
																																	yield (file_idx, batch_idx), self._cast_table(lowerCAmelCase__)
																					elif self.config.sample_by == "document":
																											SCREAMING_SNAKE_CASE_:				List[Any]							= f.read()
																											SCREAMING_SNAKE_CASE_:				Optional[int]							= pa.Table.from_arrays([pa.array([text])]     ,		names=lowerCAmelCase__)
																											yield file_idx, self._cast_table(lowerCAmelCase__)






