




"""simple docstring"""




import argparse
import dataclasses
import json
import logging
import os
import shutil
from typing import List, Optional

import datasets
from accelerate import Accelerator
from datasets import load_dataset
from finetuning import finetune
from tqdm.auto import tqdm

import transformers
from transformers import AutoConfig, set_seed
from transformers.trainer_utils import IntervalStrategy


UpperCamelCase__         =      logging.getLogger(__name__)

UpperCamelCase__         =      """pytorch_model.bin"""
@dataclasses.dataclass
class       a__   :
						_a							:   Optional[int] =   dataclasses.field(
						    metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models."""}				)
						_a							:   Optional[int] =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co."""}			,    )






@dataclasses.dataclass
class       a__   :
						_a							:   Union[str, Any] =   dataclasses.field(metadata={"""help""": """A csv or a json file containing the training data."""}				)
						_a							:   Dict =   dataclasses.field(metadata={"""help""": """A csv or a json file containing the data to predict on."""}				)
						_a							:   Optional[Any] =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """A csv or a json file containing the validation data."""}				)
						_a							:   Tuple =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """The name of the task to train on."""}			,    )
						_a							:   Optional[int] =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """The list of labels for the task."""}				)






@dataclasses.dataclass
class       a__   :
						_a							:   Optional[Any] =   dataclasses.field(
						    metadata={"""help""": """The output directory where the model predictions and checkpoints will be written."""}				)
						_a							:   Optional[Any] =   dataclasses.field(
						    default="""accuracy"""			,    metadata={"""help""": """The evaluation metric used for the task."""}				)
						_a							:   Dict =   dataclasses.field(
						    default="""no"""			,    metadata={
						        """help""": """The evaluation strategy to adopt during training. Possible values are: [\"no\", \"step\", \"epoch]"""
						    }			,    )
						_a							:   Optional[int] =   dataclasses.field(
						    default=1_0			,    metadata={"""help""": """Number of evaluation calls with no improvement after which training will be stopped."""}			,    )
						_a							:   List[Any] =   dataclasses.field(
						    default=0.0			,    metadata={
						        """help""": """How much the specified evaluation metric must improve to satisfy early stopping conditions."""
						    }			,    )
						_a							:   List[str] =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """Whether to filter the pseudo-labeled data based on the confidence score."""}			,    )
						_a							:   Optional[Any] =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """Whether to filter the pseudo-labeled data based on the validation performance."""}			,    )
						_a							:   Tuple =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """Whether to fine-tune on labeled data after pseudo training."""}			,    )
						_a							:   List[Any] =   dataclasses.field(
						    default=0.0			,    metadata={"""help""": """Confidence threshold for pseudo-labeled data filtering."""}			,    )
						_a							:   Dict =   dataclasses.field(
						    default=1_0_0			,    metadata={"""help""": """Number of evaluation calls with no improvement after which training will be stopped."""}			,    )
						_a							:   Optional[int] =   dataclasses.field(
						    default=a__			,    metadata={"""help""": """Random seed for initialization."""}			,    )







def      _a				(						SCREAMING_SNAKE_CASE_							:    Any				,				SCREAMING_SNAKE_CASE_							:    Tuple				,				SCREAMING_SNAKE_CASE_							:    Any				,				SCREAMING_SNAKE_CASE_							:    Union[str, Any]				,				SCREAMING_SNAKE_CASE_							:    List[str]				,				SCREAMING_SNAKE_CASE_							:    Optional[Any]					):
					__lowerCAmelCase           =	datasets.concatenate_datasets([infer_input, infer_output]				,				axis=1					)

					if args.do_filter_by_confidence:
										__lowerCAmelCase           =	dataset.filter(lambda SCREAMING_SNAKE_CASE_					: example["probability"] > args.confidence_threshold					)

					if args.do_filter_by_val_performance:
										assert eval_result >= 0.0 and eval_result <= 1.0
										__lowerCAmelCase           =	int(eval_result * len(_lowerCamelCase					)					)
										print(_lowerCamelCase					)
										__lowerCAmelCase           =	dataset.sort("probability"				,				reverse=_lowerCamelCase					)
										__lowerCAmelCase           =	dataset.select(range(_lowerCamelCase					)					)

					__lowerCAmelCase           =	dataset.remove_columns(["label", "probability"]					)
					__lowerCAmelCase           =	dataset.rename_column("prediction"				,				"label"					)
					__lowerCAmelCase           =	dataset.map(lambda SCREAMING_SNAKE_CASE_					: {"label": idalabel[example["label"]]}					)
					__lowerCAmelCase           =	dataset.shuffle(seed=args.seed					)

					__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				F"""train_pseudo.{args.data_file_extension}"""					)
					if args.data_file_extension == "csv":
										dataset.to_csv(_lowerCamelCase				,				index=_lowerCamelCase					)
					else:
										dataset.to_json(_lowerCamelCase					)







def      _a				(						SCREAMING_SNAKE_CASE_							:    str				,				SCREAMING_SNAKE_CASE_							:    Dict				,				SCREAMING_SNAKE_CASE_							:    Tuple				,				SCREAMING_SNAKE_CASE_							:    Optional[int]				,				**SCREAMING_SNAKE_CASE_							:    Any					):
					__lowerCAmelCase           =	Accelerator()
					# Make one log on every process with the configuration for debugging.
					logging.basicConfig(
					    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"				,				datefmt="%m/%d/%Y %H:%M:%S"				,				level=logging.INFO				,				)
					logger.info(accelerator.state					)

					# Setup logging, we only want one process per machine to log things on the
					# screen. accelerator.is_local_main_process is only True for one process per
					# machine.
					logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR					)

					if accelerator.is_local_main_process:
										datasets.utils.logging.set_verbosity_warning()
										transformers.utils.logging.set_verbosity_info()
					else:
										datasets.utils.logging.set_verbosity_error()
										transformers.utils.logging.set_verbosity_error()

					__lowerCAmelCase           =	STModelArguments(model_name_or_path=_lowerCamelCase					)
					__lowerCAmelCase           =	STDataArguments(train_file=_lowerCamelCase				,				infer_file=_lowerCamelCase					)
					__lowerCAmelCase           =	STTrainingArguments(output_dir=_lowerCamelCase					)
					__lowerCAmelCase           =	argparse.Namespace()

					for arg_class in (model_args, data_args, training_args):
										for key, value in vars(_lowerCamelCase					).items():
															setattr(_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase					)

					for key, value in kwargs.items():
										if hasattr(_lowerCamelCase				,				_lowerCamelCase					):
															setattr(_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase					)

    # Sanity checks
					__lowerCAmelCase           =	{}
					__lowerCAmelCase           =	None

					# You need to provide the training data and the data to predict on
					assert args.train_file is not None
					assert args.infer_file is not None
					__lowerCAmelCase           =	args.train_file
					__lowerCAmelCase           =	args.infer_file

					if args.evaluation_strategy != IntervalStrategy.NO.value:
										assert args.eval_file is not None
										__lowerCAmelCase           =	args.eval_file

					for key in data_files:
										__lowerCAmelCase           =	data_files[key].split("."					)[-1]
										assert extension in ["csv", "json"], F"""`{key}_file` should be a csv or a json file."""
										if args.data_file_extension is None:
															__lowerCAmelCase           =	extension
										else:
															assert extension == args.data_file_extension, F"""`{key}_file` should be a {args.data_file_extension} file`."""

					assert (
					    args.eval_metric in datasets.list_metrics()
					), F"""{args.eval_metric} not in the list of supported metrics {datasets.list_metrics()}."""

					# If passed along, set the training seed now.
					if args.seed is not None:
										set_seed(args.seed					)

					logger.info("Creating the initial data directory for self-training..."					)
					__lowerCAmelCase           =	F"""{args.output_dir}/self-train_iter-{{}}""".format
					__lowerCAmelCase           =	data_dir_format(0					)

					if accelerator.is_main_process:
										if args.output_dir is not None:
															os.makedirs(args.output_dir				,				exist_ok=_lowerCamelCase					)
															os.makedirs(_lowerCamelCase				,				exist_ok=_lowerCamelCase					)
					accelerator.wait_for_everyone()

					__lowerCAmelCase           =	None
					__lowerCAmelCase           =	None
					__lowerCAmelCase           =	0
					__lowerCAmelCase           =	False
					# Show the progress bar
					__lowerCAmelCase           =	tqdm(range(args.max_selftrain_iterations					)				,				disable=not accelerator.is_local_main_process					)

					# Self-train
					for iteration in range(0				,				int(args.max_selftrain_iterations					)					):
										__lowerCAmelCase           =	data_dir_format(_lowerCamelCase					)
										assert os.path.exists(_lowerCamelCase					)

										# Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
										# iteration > 0
										__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"stage-1"					)
										__lowerCAmelCase           =	{
										    "accelerator": accelerator,
										    "model_name_or_path": args.model_name_or_path,
										    "cache_dir": args.cache_dir,
										    "do_train": True,
										    "train_file": data_files["train"] if iteration == 0 else data_files["train_pseudo"],
										    "do_eval": True if args.eval_file is not None else False,
										    "eval_file": data_files["eval"],
										    "do_predict": True,
										    "infer_file": data_files["infer"],
										    "task_name": args.task_name,
										    "label_list": args.label_list,
										    "output_dir": current_output_dir,
										    "eval_metric": args.eval_metric,
										    "evaluation_strategy": args.evaluation_strategy,
										    "early_stopping_patience": args.early_stopping_patience,
										    "early_stopping_threshold": args.early_stopping_threshold,
										    "seed": args.seed,
										}
										# Add additional training arguments
										for key, value in kwargs.items():
															if key not in arguments_dict and not hasattr(_lowerCamelCase				,				_lowerCamelCase					):
																				arguments_dict.update({key: value}					)

										__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"best-checkpoint"				,				_lowerCamelCase					)
										if os.path.exists(_lowerCamelCase					):
															logger.info(
															    "Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 1."				,				_lowerCamelCase				,				_lowerCamelCase				,				)
										else:
															logger.info("***** Running self-training: iteration: %d, stage: 1 *****"				,				_lowerCamelCase					)
															finetune(**_lowerCamelCase					)
															accelerator.wait_for_everyone()
															assert os.path.exists(_lowerCamelCase					)
															logger.info("Self-training job completed: iteration: %d, stage: 1."				,				_lowerCamelCase					)

										if iteration > 0 and args.finetune_on_labeled_data:
															# Stage 2 (optional): fine-tuning on the original labeled data
															__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"best-checkpoint"					)
															__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"stage-2"					)
															# Update arguments_dict
															__lowerCAmelCase           =	model_path
															__lowerCAmelCase           =	data_files["train"]
															__lowerCAmelCase           =	current_output_dir

															__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"best-checkpoint"				,				_lowerCamelCase					)
															if os.path.exists(_lowerCamelCase					):
																				logger.info(
																				    "Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 2."				,				_lowerCamelCase				,				_lowerCamelCase				,				)
															else:
																				logger.info("***** Running self-training: iteration: %d, stage: 2 *****"				,				_lowerCamelCase					)
																				finetune(**_lowerCamelCase					)
																				accelerator.wait_for_everyone()
																				assert os.path.exists(_lowerCamelCase					)
																				logger.info("Self-training job completed: iteration: %d, stage: 2."				,				_lowerCamelCase					)

										__lowerCAmelCase           =	iteration
										__lowerCAmelCase           =	data_dir_format(iteration + 1					)

										__lowerCAmelCase           =	AutoConfig.from_pretrained(os.path.join(_lowerCamelCase				,				"best-checkpoint"					)					)
										__lowerCAmelCase           =	config.idalabel
										__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"eval_results_best-checkpoint.json"					)
										__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"test_results_best-checkpoint.json"					)
										assert os.path.exists(_lowerCamelCase					)

										with open(_lowerCamelCase				,				"r"					) as f:
															__lowerCAmelCase           =	float(json.load(_lowerCamelCase					)[args.eval_metric]					)
										__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				"infer_output_best-checkpoint.csv"					)
										assert os.path.exists(_lowerCamelCase					)
										# Loading the dataset from local csv or json files.
										__lowerCAmelCase           =	load_dataset(args.data_file_extension				,				data_files={"data": data_files["infer"]}					)["data"]
										__lowerCAmelCase           =	load_dataset("csv"				,				data_files={"data": infer_output_file}					)["data"]

										if accelerator.is_main_process:
															os.makedirs(_lowerCamelCase				,				exist_ok=_lowerCamelCase					)
															shutil.copy(_lowerCamelCase				,				os.path.join(_lowerCamelCase				,				F"""eval_results_iter-{iteration}.json"""					)					)
															if os.path.exists(_lowerCamelCase					):
																				shutil.copy(_lowerCamelCase				,				os.path.join(_lowerCamelCase				,				F"""test_results_iter-{iteration}.json"""					)					)
															create_pseudo_labeled_data(_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase					)
										accelerator.wait_for_everyone()

										__lowerCAmelCase           =	os.path.join(_lowerCamelCase				,				F"""train_pseudo.{args.data_file_extension}"""					)

										if args.evaluation_strategy != IntervalStrategy.NO.value:
															__lowerCAmelCase           =	eval_result

															if best_iteration is None:
																				__lowerCAmelCase           =	new_iteration
																				__lowerCAmelCase           =	new_eval_result
															else:
																				if new_eval_result - best_eval_result > args.early_stopping_threshold:
																									__lowerCAmelCase           =	new_iteration
																									__lowerCAmelCase           =	new_eval_result
																									__lowerCAmelCase           =	0
																				else:
																									if new_eval_result == best_eval_result:
																														__lowerCAmelCase           =	new_iteration
																														__lowerCAmelCase           =	new_eval_result
																									early_stopping_patience_counter += 1

																				if early_stopping_patience_counter >= args.early_stopping_patience:
																									__lowerCAmelCase           =	True

										progress_bar.update(1					)

										if should_training_stop:
															break

					if best_iteration is not None:
										# Save the best iteration
										logger.info("Best iteration: %d"				,				_lowerCamelCase					)
										logger.info("Best evaluation result: %s = %f"				,				args.eval_metric				,				_lowerCamelCase					)
										accelerator.wait_for_everyone()
										if accelerator.is_main_process:
															shutil.copy(
															    os.path.join(_lowerCamelCase				,				F"""eval_results_iter-{iteration}.json"""					)				,				os.path.join(_lowerCamelCase				,				"eval_results_best-iteration.json"					)				,				)
					else:
										# Assume that the last iteration is the best
										logger.info("Best iteration: %d"				,				args.max_selftrain_iterations - 1					)
										logger.info("Best evaluation result: %s = %f"				,				args.eval_metric				,				_lowerCamelCase					)
										accelerator.wait_for_everyone()
										if accelerator.is_main_process:
															shutil.copy(
															    os.path.join(_lowerCamelCase				,				F"""eval_results_iter-{args.max_selftrain_iterations - 1}.json"""					)				,				os.path.join(_lowerCamelCase				,				"eval_results_best-iteration.json"					)				,				)










import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


UpperCamelCase__         =      logging.getLogger(__name__)
class       a__   :






						def __init__(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											if not self.initialized:
																__lowerCAmelCase           =	RagRetriever(
																    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      index=_A		,      init_retrieval=_A		,      )
																__lowerCAmelCase           =	True






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.retriever.index.init_index()






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.retriever._main_retrieve(_A		,      _A      )
											return doc_ids, retrieved_doc_embeds






class       a__   ( snake_case__				):






						def __init__(       self		,      _A		,      _A		,      _A		,      _A		,      _A=None      ):




											"""simple docstring"""



											if index is not None and index.is_initialized() and len(_A      ) > 0:
																raise ValueError(
																    "When using Ray for distributed fine-tuning, "
																    "you'll need to provide the paths instead, "
																    "as the dataset and the index are loaded "
																    "separately. More info in examples/rag/use_own_knowledge_dataset.py "      )
											super().__init__(
											    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      index=_A		,      init_retrieval=_A		,      )
											__lowerCAmelCase           =	retrieval_workers
											if len(self.retrieval_workers      ) > 0:
																ray.get(
																    [
																        worker.create_rag_retriever.remote(_A		,      _A		,      _A		,      _A      )
																        for worker in self.retrieval_workers
																    ]      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											logger.info("initializing retrieval"      )

											if len(self.retrieval_workers      ) > 0:
																ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]      )
											else:
																# Non-distributed training. Load index into this same process.
																self.index.init_index()






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											if len(self.retrieval_workers      ) > 0:
																# Select a random retrieval actor.
																__lowerCAmelCase           =	self.retrieval_workers[random.randint(0		,      len(self.retrieval_workers      ) - 1      )]
																__lowerCAmelCase       ,							__lowerCAmelCase           =	ray.get(random_worker.retrieve.remote(_A		,      _A      )      )
											else:
																__lowerCAmelCase       ,							__lowerCAmelCase           =	self._main_retrieve(_A		,      _A      )
											return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(_A      )






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls		,      _A		,      _A=None		,      **_A      ):




											"""simple docstring"""



											return super(_A		,      cls      ).get_tokenizers(_A		,      _A		,      **_A      )






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls		,      _A		,      _A		,      _A=None		,      **_A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	kwargs.pop("config"		,      _A      ) or RagConfig.from_pretrained(_A		,      **_A      )
											__lowerCAmelCase           =	RagTokenizer.from_pretrained(_A		,      config=_A      )
											__lowerCAmelCase           =	rag_tokenizer.question_encoder
											__lowerCAmelCase           =	rag_tokenizer.generator
											if indexed_dataset is not None:
																__lowerCAmelCase           =	"custom"
																__lowerCAmelCase           =	CustomHFIndex(config.retrieval_vector_size		,      _A      )
											else:
																__lowerCAmelCase           =	cls._build_index(_A      )
											return cls(
											    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      retrieval_workers=_A		,      index=_A		,      )




