




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    __lowercase						:
   _a							=      dataclasses.field(
       metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models."""} )
   _a							=      dataclasses.field(
       default=__magic_name__       ,							metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co."""}       ,							)





@dataclasses.dataclass
class    __lowercase						:
   _a							=      dataclasses.field(metadata={"""help""": """A csv or a json file containing the training data."""} )
   _a							=      dataclasses.field(metadata={"""help""": """A csv or a json file containing the data to predict on."""} )
   _a							=      dataclasses.field(
       default=__magic_name__       ,							metadata={"""help""": """A csv or a json file containing the validation data."""} )
   _a							=      dataclasses.field(
       default=__magic_name__       ,							metadata={"""help""": """The name of the task to train on."""}       ,							)
   _a							=      dataclasses.field(
       default=__magic_name__       ,							metadata={"""help""": """The list of labels for the task."""} )





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


def 						SCREAMING_SNAKE_CASE       (      a_	:							str				,					a_	:							int				,					a_	:							List[Any]				,					a_	:							List[Any]				,					a_	:							int				,					a_	:							Optional[int]			):
  __a			=					datasets.concatenate_datasets([infer_input, infer_output]				,					axis=1			)

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

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

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

  __a			=					os.path.join(a_				,					f"train_pseudo.{args.data_file_extension}"			)
  if args.data_file_extension == "csv":
    dataset.to_csv(a_				,					index=a_			)
  else:
    dataset.to_json(a_			)


def 						SCREAMING_SNAKE_CASE       (      a_	:							int				,					a_	:							Dict				,					a_	:							Any				,					a_	:							List[str]				,					**a_	:							Union[str, Any]			):
  __a			=					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()

  __a			=					STModelArguments(model_name_or_path=a_			)
  __a			=					STDataArguments(train_file=a_				,					infer_file=a_			)
  __a			=					STTrainingArguments(output_dir=a_			)
  __a			=					argparse.Namespace()

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

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

    # Sanity checks
  __a			=					{}
  __a			=					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
  __a			=					args.train_file
  __a			=					args.infer_file

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

  for key in data_files:
    __a			=					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:
      __a			=					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...'			)
  __a			=					f"{args.output_dir}/self-train_iter-{{}}".format
  __a			=					data_dir_format(0			)

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

  __a			=					None
  __a			=					None
  __a			=					0
  __a			=					False
  # Show the progress bar
  __a			=					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			)			):
    __a			=					data_dir_format(a_			)
    assert os.path.exists(a_			)

    # Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
    # iteration > 0
    __a			=					os.path.join(a_				,					'stage-1'			)
    __a			=					{
        '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(a_				,					a_			):
        arguments_dict.update({key: value}			)

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

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

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

    __a			=					iteration
    __a			=					data_dir_format(iteration + 1			)

    __a			=					AutoConfig.from_pretrained(os.path.join(a_				,					'best-checkpoint'			)			)
    __a			=					config.idalabel
    __a			=					os.path.join(a_				,					'eval_results_best-checkpoint.json'			)
    __a			=					os.path.join(a_				,					'test_results_best-checkpoint.json'			)
    assert os.path.exists(a_			)

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

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

    __a			=					os.path.join(a_				,					f"train_pseudo.{args.data_file_extension}"			)

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

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

        if early_stopping_patience_counter >= args.early_stopping_patience:
          __a			=					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'				,					a_			)
    logger.info('Best evaluation result: %s = %f'				,					args.eval_metric				,					a_			)
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
      shutil.copy(
          os.path.join(a_				,					f"eval_results_iter-{iteration}.json"			)				,					os.path.join(a_				,					'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				,					a_			)
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
      shutil.copy(
          os.path.join(a_				,					f"eval_results_iter-{args.max_selftrain_iterations - 1}.json"			)				,					os.path.join(a_				,					'eval_results_best-iteration.json'			)				,					)









'''simple docstring'''


# Usage:
# ./gen-card-allenai-wmt16.py

import os
from pathlib import Path


def 						SCREAMING_SNAKE_CASE       (      a_	:							Optional[Any]				,					a_	:							int				,					a_	:							Tuple				,					a_	:							Tuple			):
  __a			=					{
      'en': 'Machine learning is great, isn\'t it?',
      'ru': 'Машинное обучение - это здорово, не так ли?',
      'de': 'Maschinelles Lernen ist großartig, nicht wahr?',
  }

  # BLUE scores as follows:
  # "pair": [fairseq, transformers]
  __a			=					{
      'wmt16-en-de-dist-12-1': [28.3, 27.52],
      'wmt16-en-de-dist-6-1': [27.4, 27.11],
      'wmt16-en-de-12-1': [26.9, 25.75],
  }
  __a			=					f"{src_lang}-{tgt_lang}"

  __a			=					f"\n---\nlanguage:\n- {src_lang}\n- {tgt_lang}\nthumbnail:\ntags:\n- translation\n- wmt16\n- allenai\nlicense: apache-2.0\ndatasets:\n- wmt16\nmetrics:\n- bleu\n---\n\n# FSMT\n\n## Model description\n\nThis is a ported version of fairseq-based [wmt16 transformer](https://github.com/jungokasai/deep-shallow/) for {src_lang}-{tgt_lang}.\n\nFor more details, please, see [Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation](https://arxiv.org/abs/2006.10369).\n\nAll 3 models are available:\n\n* [wmt16-en-de-dist-12-1](https://huggingface.co/allenai/wmt16-en-de-dist-12-1)\n* [wmt16-en-de-dist-6-1](https://huggingface.co/allenai/wmt16-en-de-dist-6-1)\n* [wmt16-en-de-12-1](https://huggingface.co/allenai/wmt16-en-de-12-1)\n\n\n## Intended uses & limitations\n\n#### How to use\n\n```python\nfrom transformers import FSMTForConditionalGeneration, FSMTTokenizer\nmname = \"allenai/{model_name}\"\ntokenizer = FSMTTokenizer.from_pretrained(mname)\nmodel = FSMTForConditionalGeneration.from_pretrained(mname)\n\ninput = \"{texts[src_lang]}\"\ninput_ids = tokenizer.encode(input, return_tensors=\"pt\")\noutputs = model.generate(input_ids)\ndecoded = tokenizer.decode(outputs[0], skip_special_tokens=True)\nprint(decoded) # {texts[tgt_lang]}\n\n```\n\n#### Limitations and bias\n\n\n## Training data\n\nPretrained weights were left identical to the original model released by allenai. For more details, please, see the [paper](https://arxiv.org/abs/2006.10369).\n\n## Eval results\n\nHere are the BLEU scores:\n\nmodel   | fairseq | transformers\n-------|---------|----------\n{model_name}  | {scores[model_name][0]} | {scores[model_name][1]}\n\nThe score is slightly below the score reported in the paper, as the researchers don't use `sacrebleu` and measure the score on tokenized outputs. `transformers` score was measured using `sacrebleu` on detokenized outputs.\n\nThe score was calculated using this code:\n\n```bash\ngit clone https://github.com/huggingface/transformers\ncd transformers\nexport PAIR={pair}\nexport DATA_DIR=data/$PAIR\nexport SAVE_DIR=data/$PAIR\nexport BS=8\nexport NUM_BEAMS=5\nmkdir -p $DATA_DIR\nsacrebleu -t wmt16 -l $PAIR --echo src > $DATA_DIR/val.source\nsacrebleu -t wmt16 -l $PAIR --echo ref > $DATA_DIR/val.target\necho $PAIR\nPYTHONPATH=\"src:examples/seq2seq\" python examples/seq2seq/run_eval.py allenai/{model_name} $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS\n```\n\n## Data Sources\n\n- [training, etc.](http://www.statmt.org/wmt16/)\n- [test set](http://matrix.statmt.org/test_sets/newstest2016.tgz?1504722372)\n\n\n### BibTeX entry and citation info\n\n```\n@misc{{kasai2020deep,\n    title={{Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation}},\n    author={{Jungo Kasai and Nikolaos Pappas and Hao Peng and James Cross and Noah A. Smith}},\n    year={{2020}},\n    eprint={{2006.10369}},\n    archivePrefix={{arXiv}},\n    primaryClass={{cs.CL}}\n}}\n```\n\n"
  model_card_dir.mkdir(parents=a_				,					exist_ok=a_			)
  __a			=					os.path.join(a_				,					'README.md'			)
  print(f"Generating {path}"			)
  with open(a_				,					'w'				,					encoding='utf-8'			) as f:
    f.write(a_			)


# make sure we are under the root of the project
UpperCAmelCase_ 		=      Path(__file__).resolve().parent.parent.parent
UpperCAmelCase_ 		=      repo_dir / "model_cards"

for model_name in ["wmt16-en-de-dist-12-1", "wmt16-en-de-dist-6-1", "wmt16-en-de-12-1"]:
  UpperCAmelCase_ 		=      model_cards_dir / "allenai" / model_name
  write_model_card(model_card_dir, src_lang="en", tgt_lang="de", model_name=model_name)



