


'''simple docstring'''







import argparse
import json
import logging
import os
import shutil
import sys
import tempfile
import unittest
from unittest import mock

import torch
from accelerate.utils import write_basic_config

from transformers.testing_utils import TestCasePlus, get_gpu_count, run_command, slow, torch_device
from transformers.utils import is_apex_available


logging.basicConfig(level=logging.DEBUG)

a__  :					str        =       logging.getLogger()







def _UpperCamelCase       (     )							->		Union[str, Any]:


      '''simple docstring'''



      UpperCamelCase__										=				argparse.ArgumentParser()
      parser.add_argument("-f"    )
      UpperCamelCase__										=				parser.parse_args()
      return args.f







def _UpperCamelCase       (     __A    )							->		Optional[Any]:


      '''simple docstring'''



      UpperCamelCase__										=				{}
      UpperCamelCase__										=				os.path.join(SCREAMING_SNAKE_CASE__      ,			"all_results.json"    )
      if os.path.exists(SCREAMING_SNAKE_CASE__    ):
            with open(SCREAMING_SNAKE_CASE__      ,			"r"    ) as f:
                  UpperCamelCase__										=				json.load(SCREAMING_SNAKE_CASE__    )
      else:
            raise ValueError(F'''can\'t find {path}'''    )
      return results







def _UpperCamelCase       (     )							->		str:


      '''simple docstring'''



      UpperCamelCase__										=				torch.cuda.is_available() and torch_device == """cuda"""
      return is_using_cuda and is_apex_available()


a__  :					List[str]        =       logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)



class 					lowercase_     (    _snake_case   ):


    @classmethod
    def    __a							(     cls    ):
          # Write Accelerate config, will pick up on CPU, GPU, and multi-GPU
          UpperCamelCase__										=				tempfile.mkdtemp()
          UpperCamelCase__										=				os.path.join(cls.tmpdir  ,					"default_config.yml"    )
          write_basic_config(save_location=cls.configPath    )
          UpperCamelCase__										=				["""accelerate""", """launch""", """--config_file""", cls.configPath]


    @classmethod
    def    __a							(     cls    ):
          shutil.rmtree(cls.tmpdir    )


    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/text-classification/run_glue_no_trainer.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --seed=42
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          if is_cuda_and_apex_available():
                testargs.append("--fp16"    )

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertGreaterEqual(result["eval_accuracy"]  ,					0.75    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"glue_no_trainer"    )    )    )


    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/language-modeling/run_clm_no_trainer.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --block_size 128
            --per_device_train_batch_size 5
            --per_device_eval_batch_size 5
            --num_train_epochs 2
            --output_dir {tmp_dir}
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          if torch.cuda.device_count() > 1:
                # Skipping because there are not enough batches to train the model + would need a drop_last to work.
                return

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertLess(result["perplexity"]  ,					1_00    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"clm_no_trainer"    )    )    )


    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/language-modeling/run_mlm_no_trainer.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --output_dir {tmp_dir}
            --num_train_epochs=1
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertLess(result["perplexity"]  ,					42    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"mlm_no_trainer"    )    )    )


    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          # with so little data distributed training needs more epochs to get the score on par with 0/1 gpu
          UpperCamelCase__										=				7 if get_gpu_count() > 1 else 2

          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/token-classification/run_ner_no_trainer.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertGreaterEqual(result["eval_accuracy"]  ,					0.75    )
          self.assertLess(result["train_loss"]  ,					0.5    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"ner_no_trainer"    )    )    )


    @unittest.skip(reason="Fix me @muellerzr"    )
    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/question-answering/run_qa_no_trainer.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --seed=42
            --max_train_steps=10
            --num_warmup_steps=2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          # Because we use --version_2_with_negative the testing script uses SQuAD v2 metrics.
          self.assertGreaterEqual(result["eval_f1"]  ,					28    )
          self.assertGreaterEqual(result["eval_exact"]  ,					28    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"qa_no_trainer"    )    )    )


    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/multiple-choice/run_swag_no_trainer.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/swag/sample.json
            --validation_file tests/fixtures/tests_samples/swag/sample.json
            --output_dir {tmp_dir}
            --max_train_steps=20
            --num_warmup_steps=2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --with_tracking
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertGreaterEqual(result["eval_accuracy"]  ,					0.8    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"swag_no_trainer"    )    )    )


    @slow
    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/summarization/run_summarization_no_trainer.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --max_train_steps=50
            --num_warmup_steps=8
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertGreaterEqual(result["eval_rouge1"]  ,					10    )
          self.assertGreaterEqual(result["eval_rouge2"]  ,					2    )
          self.assertGreaterEqual(result["eval_rougeL"]  ,					7    )
          self.assertGreaterEqual(result["eval_rougeLsum"]  ,					7    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"summarization_no_trainer"    )    )    )


    @slow
    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/translation/run_translation_no_trainer.py
            --model_name_or_path sshleifer/student_marian_en_ro_6_1
            --source_lang en
            --target_lang ro
            --train_file tests/fixtures/tests_samples/wmt16/sample.json
            --validation_file tests/fixtures/tests_samples/wmt16/sample.json
            --output_dir {tmp_dir}
            --max_train_steps=50
            --num_warmup_steps=8
            --num_beams=6
            --learning_rate=3e-3
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --source_lang en_XX
            --target_lang ro_RO
            --checkpointing_steps epoch
            --with_tracking
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertGreaterEqual(result["eval_bleu"]  ,					30    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"epoch_0"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"translation_no_trainer"    )    )    )


    @slow
    def    __a							(     self    ):
          UpperCamelCase__										=				logging.StreamHandler(sys.stdout    )
          logger.addHandler(__snake_case    )

          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
            --dataset_name huggingface/semantic-segmentation-test-sample
            --output_dir {tmp_dir}
            --max_train_steps=10
            --num_warmup_steps=2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --checkpointing_steps epoch
        '''.split()

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          self.assertGreaterEqual(result["eval_overall_accuracy"]  ,					0.10    )


    @mock.patch.dict(os.environ  ,					{"WANDB_MODE": "offline"}    )
    def    __a							(     self    ):
          UpperCamelCase__										=				self.get_auto_remove_tmp_dir()
          UpperCamelCase__										=				f'''
            {self.examples_dir}/pytorch/image-classification/run_image_classification_no_trainer.py
            --model_name_or_path google/vit-base-patch16-224-in21k
            --dataset_name hf-internal-testing/cats_vs_dogs_sample
            --learning_rate 1e-4
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 1
            --max_train_steps 2
            --train_val_split 0.1
            --seed 42
            --output_dir {tmp_dir}
            --with_tracking
            --checkpointing_steps 1
        '''.split()

          if is_cuda_and_apex_available():
                testargs.append("--fp16"    )

          run_command(self._launch_args + testargs    )
          UpperCamelCase__										=				get_results(__snake_case    )
          # The base model scores a 25%
          self.assertGreaterEqual(result["eval_accuracy"]  ,					0.6    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"step_1"    )    )    )
          self.assertTrue(os.path.exists(os.path.join(__snake_case  ,					"image_classification_no_trainer"    )    )    )









'''simple docstring'''


from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map







def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       Tuple     		=      []
				_SCREAMING_SNAKE_CASE :       Dict     		=      []
				_SCREAMING_SNAKE_CASE :       str     		=      []

				for rt in rc.restypes:
								_SCREAMING_SNAKE_CASE :       Optional[int]     		=      rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
								restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]     )
								_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      {name: i for i, name in enumerate(SCREAMING_SNAKE_CASE__     )}
								restype_atomaa_to_atomaa_list.append(
								    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]     )

								restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]     )

				# Add dummy mapping for restype 'UNK'
				restype_atomaa_to_atomaa_list.append([0] * 14     )
				restype_atomaa_to_atomaa_list.append([0] * 37     )
				restype_atomaa_mask_list.append([0.0] * 14     )

				_SCREAMING_SNAKE_CASE :       Dict     		=      torch.tensor(
				    SCREAMING_SNAKE_CASE__							,      dtype=torch.intaa							,      device=protein["""aatype"""].device							,      )
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      torch.tensor(
				    SCREAMING_SNAKE_CASE__							,      dtype=torch.intaa							,      device=protein["""aatype"""].device							,      )
				_SCREAMING_SNAKE_CASE :       str     		=      torch.tensor(
				    SCREAMING_SNAKE_CASE__							,      dtype=torch.floataa							,      device=protein["""aatype"""].device							,      )
				_SCREAMING_SNAKE_CASE :       int     		=      protein["""aatype"""].to(torch.long     )

				# create the mapping for (residx, atom14) --> atom37, i.e. an array
				# with shape (num_res, 14) containing the atom37 indices for this protein
				_SCREAMING_SNAKE_CASE :       List[str]     		=      restype_atomaa_to_atomaa[protein_aatype]
				_SCREAMING_SNAKE_CASE :       str     		=      restype_atomaa_mask[protein_aatype]

				_SCREAMING_SNAKE_CASE :       List[Any]     		=      residx_atomaa_mask
				_SCREAMING_SNAKE_CASE :       List[str]     		=      residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				_SCREAMING_SNAKE_CASE :       int     		=      restype_atomaa_to_atomaa[protein_aatype]
				_SCREAMING_SNAKE_CASE :       Dict     		=      residx_atomaa_to_atomaa.long()

				# create the corresponding mask
				_SCREAMING_SNAKE_CASE :       str     		=      torch.zeros([21, 37]							,      dtype=torch.floataa							,      device=protein["""aatype"""].device     )
				for restype, restype_letter in enumerate(rc.restypes     ):
								_SCREAMING_SNAKE_CASE :       int     		=      rc.restype_atoa[restype_letter]
								_SCREAMING_SNAKE_CASE :       Dict     		=      rc.residue_atoms[restype_name]
								for atom_name in atom_names:
												_SCREAMING_SNAKE_CASE :       List[Any]     		=      rc.atom_order[atom_name]
												_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      1

				_SCREAMING_SNAKE_CASE :       int     		=      restype_atomaa_mask[protein_aatype]
				_SCREAMING_SNAKE_CASE :       int     		=      residx_atomaa_mask

				return protein







def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      tree_map(lambda SCREAMING_SNAKE_CASE__     : torch.tensor(SCREAMING_SNAKE_CASE__							,      device=batch["""aatype"""].device     )							,      SCREAMING_SNAKE_CASE__							,      np.ndarray     )
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      tensor_tree_map(lambda SCREAMING_SNAKE_CASE__     : np.array(SCREAMING_SNAKE_CASE__     )							,      make_atomaa_masks(SCREAMING_SNAKE_CASE__     )     )
				return out




