




import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
		from sagemaker import Session, TrainingJobAnalytics
		from sagemaker.huggingface import HuggingFace



@pytest.mark.skipif(
    literal_eval(os.getenv('''TEST_SAGEMAKER'''			,  '''False'''      )      ) is not True			,  reason='''Skipping test because should only be run when releasing minor transformers version'''			,  )
@pytest.mark.usefixtures('''sm_env'''      )
@parameterized_class(
    [
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_glue.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_5_0, '''eval_accuracy''': 0.7, '''eval_loss''': 0.6},
        },
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_ddp.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.7, '''eval_loss''': 0.6},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf_dist.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.6, '''eval_loss''': 0.7},
        },
    ]      )
class   lowerCamelCase				(unittest.TestCase      ):







		'''simple docstring'''



		def  __UpperCAmelCase      (							self       )     ->							Any:
							if self.framework == "pytorch":
												subprocess.run(
												    f"cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py".split()    ,		encoding='utf-8'    ,		check=UpperCAmelCase_    ,		)
							assert hasattr(self    ,		'env'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Tuple:
							UpperCAmelCase_   :				Union[str, Any]						=						f"{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"
							# distributed data settings
							UpperCAmelCase_   :				Optional[int]						=						{"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

							# creates estimator
							return HuggingFace(
							    entry_point=self.script    ,		source_dir=self.env.test_path    ,		role=self.env.role    ,		image_uri=self.env.image_uri    ,		base_job_name=UpperCAmelCase_    ,		instance_count=UpperCAmelCase_    ,		instance_type=self.instance_type    ,		debugger_hook_config=UpperCAmelCase_    ,		hyperparameters={**self.env.distributed_hyperparameters, 'model_name_or_path': self.model_name_or_path}    ,		metric_definitions=self.env.metric_definitions    ,		distribution=UpperCAmelCase_    ,		py_version='py36'    ,		)



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							str:
							TrainingJobAnalytics(UpperCAmelCase_       ).export_csv(f"{self.env.test_path}/{job_name}_metrics.csv"       )



		@parameterized.expand([(2,)]       )
		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Dict:
							UpperCAmelCase_   :				List[Any]						=						self.create_estimator(UpperCAmelCase_       )

							# run training
							estimator.fit()

							# result dataframe
							UpperCAmelCase_   :				Optional[Any]						=						TrainingJobAnalytics(estimator.latest_training_job.name       ).dataframe()

							# extract kpis
							UpperCAmelCase_   :				str						=						list(result_metrics_df[result_metrics_df.metric_name == 'eval_accuracy']['value']       )
							UpperCAmelCase_   :				Tuple						=						list(result_metrics_df[result_metrics_df.metric_name == 'eval_loss']['value']       )
							# get train time from SageMaker job, this includes starting, preprocessing, stopping
							UpperCAmelCase_   :				str						=						(
							    Session().describe_training_job(estimator.latest_training_job.name       ).get('TrainingTimeInSeconds'    ,		9_9_9_9_9_9       )
							)

							# assert kpis
							assert train_runtime <= self.results["train_runtime"]
							assert all(t >= self.results['eval_accuracy'] for t in eval_accuracy       )
							assert all(t <= self.results['eval_loss'] for t in eval_loss       )

							# dump tests result into json file to share in PR
							with open(f"{estimator.latest_training_job.name}.json"    ,		'w'       ) as outfile:
												json.dump({'train_time': train_runtime, 'eval_accuracy': eval_accuracy, 'eval_loss': eval_loss}    ,		UpperCAmelCase_       )














from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from ..image_utils import load_image


if is_torch_available():
		import torch

		from ..models.auto.modeling_auto import MODEL_FOR_OBJECT_DETECTION_MAPPING, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING

__UpperCAmelCase     							=						logging.get_logger(__name__)


__UpperCAmelCase     							=						Dict[str, Any]
__UpperCAmelCase     							=						List[Prediction]



@add_end_docstrings(_snake_case      )
class   lowerCamelCase				(_snake_case      ):







		'''simple docstring'''



		def __init__(							self    ,		*_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[int]:
							super().__init__(*_UpperCamelCase    ,		**_UpperCamelCase       )

							if self.framework == "tf":
												raise ValueError(f"The {self.__class__} is only available in PyTorch."       )

							requires_backends(self    ,		'vision'       )
							self.check_model_type(
							    dict(MODEL_FOR_OBJECT_DETECTION_MAPPING.items() + MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items()       )       )



		def  __UpperCAmelCase      (							self    ,		**_UpperCamelCase       )     ->							Tuple:
							UpperCAmelCase_   :				str						=						{}
							if "threshold" in kwargs:
												UpperCAmelCase_   :				Tuple						=						kwargs['threshold']
							return {}, {}, postprocess_kwargs



		def __call__(							self    ,		*_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Union[Predictions, List[Prediction]]:
							return super().__call__(*_UpperCamelCase    ,		**_UpperCamelCase       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Dict:
							UpperCAmelCase_   :				Any						=						load_image(_UpperCamelCase       )
							UpperCAmelCase_   :				Any						=						torch.IntTensor([[image.height, image.width]]       )
							UpperCAmelCase_   :				str						=						self.image_processor(images=[image]    ,		return_tensors='pt'       )
							if self.tokenizer is not None:
												UpperCAmelCase_   :				List[str]						=						self.tokenizer(text=inputs['words']    ,		boxes=inputs['boxes']    ,		return_tensors='pt'       )
							UpperCAmelCase_   :				Any						=						target_size
							return inputs



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							List[Any]:
							UpperCAmelCase_   :				Any						=						model_inputs.pop('target_size'       )
							UpperCAmelCase_   :				Optional[Any]						=						self.model(**_UpperCamelCase       )
							UpperCAmelCase_   :				Dict						=						outputs.__class__({'target_size': target_size, **outputs}       )
							if self.tokenizer is not None:
												UpperCAmelCase_   :				List[str]						=						model_inputs['bbox']
							return model_outputs



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase=0.9       )     ->							List[str]:
							UpperCAmelCase_   :				List[Any]						=						model_outputs['target_size']
							if self.tokenizer is not None:
												# This is a LayoutLMForTokenClassification variant.
												# The OCR got the boxes and the model classified the words.
												UpperCAmelCase_	,  UpperCAmelCase_   :				str						=						target_size[0].tolist()

												def unnormalize(_UpperCamelCase       ):
																	return self._get_bounding_box(
																	    torch.Tensor(
																	        [
																	            (width * bbox[0] / 1_0_0_0),
																	            (height * bbox[1] / 1_0_0_0),
																	            (width * bbox[2] / 1_0_0_0),
																	            (height * bbox[3] / 1_0_0_0),
																	        ]       )       )

												UpperCAmelCase_	,  UpperCAmelCase_   :				List[Any]						=						model_outputs['logits'].squeeze(0       ).softmax(dim=-1       ).max(dim=-1       )
												UpperCAmelCase_   :				Union[str, Any]						=						[self.model.config.idalabel[prediction] for prediction in classes.tolist()]
												UpperCAmelCase_   :				Tuple						=						[unnormalize(_UpperCamelCase       ) for bbox in model_outputs['bbox'].squeeze(0       )]
												UpperCAmelCase_   :				List[str]						=						['score', 'label', 'box']
												UpperCAmelCase_   :				Any						=						[dict(zip(_UpperCamelCase    ,		_UpperCamelCase       )       ) for vals in zip(scores.tolist()    ,		_UpperCamelCase    ,		_UpperCamelCase       ) if vals[0] > threshold]
							else:
												# This is a regular ForObjectDetectionModel
												UpperCAmelCase_   :				Union[str, Any]						=						self.image_processor.post_process_object_detection(_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )
												UpperCAmelCase_   :				str						=						raw_annotations[0]
												UpperCAmelCase_   :				str						=						raw_annotation['scores']
												UpperCAmelCase_   :				Tuple						=						raw_annotation['labels']
												UpperCAmelCase_   :				List[Any]						=						raw_annotation['boxes']

												UpperCAmelCase_   :				Union[str, Any]						=						scores.tolist()
												UpperCAmelCase_   :				int						=						[self.model.config.idalabel[label.item()] for label in labels]
												UpperCAmelCase_   :				Any						=						[self._get_bounding_box(_UpperCamelCase       ) for box in boxes]

												# {"scores": [...], ...} --> [{"score":x, ...}, ...]
												UpperCAmelCase_   :				int						=						['score', 'label', 'box']
												UpperCAmelCase_   :				Dict						=						[
												    dict(zip(_UpperCamelCase    ,		_UpperCamelCase       )       )
												    for vals in zip(raw_annotation['scores']    ,		raw_annotation['labels']    ,		raw_annotation['boxes']       )
												]

							return annotation



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Dict[str, int]:
							if self.framework != "pt":
												raise ValueError('The ObjectDetectionPipeline is only available in PyTorch.'       )
							UpperCAmelCase_	,  UpperCAmelCase_	,  UpperCAmelCase_	,  UpperCAmelCase_   :				Dict						=						box.int().tolist()
							UpperCAmelCase_   :				str						=						{
							    'xmin': xmin,
							    'ymin': ymin,
							    'xmax': xmax,
							    'ymax': ymax,
							}
							return bbox






