



import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SegformerConfig,
    SegformerForImageClassification,
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowercase    : List[str]    	=  logging.get_logger(__name__)


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase=False					):
       _UpperCamelCase              =						OrderedDict()
       for key, value in state_dict.items():
              if encoder_only and not key.startswith('''head'''					):
                     _UpperCamelCase              =						'''segformer.encoder.''' + key
              if key.startswith('''backbone'''					):
                     _UpperCamelCase              =						key.replace('''backbone'''    ,       '''segformer.encoder'''					)
              if "patch_embed" in key:
                     # replace for example patch_embed1 by patch_embeddings.0
                     _UpperCamelCase              =						key[key.find('''patch_embed'''					) + len('''patch_embed'''					)]
                     _UpperCamelCase              =						key.replace(f'''patch_embed{idx}'''    ,       f'''patch_embeddings.{int(lowerCAmelCase					)-1}'''					)
              if "norm" in key:
                     _UpperCamelCase              =						key.replace('''norm'''    ,       '''layer_norm'''					)
              if "segformer.encoder.layer_norm" in key:
                     # replace for example layer_norm1 by layer_norm.0
                     _UpperCamelCase              =						key[key.find('''segformer.encoder.layer_norm'''					) + len('''segformer.encoder.layer_norm'''					)]
                     _UpperCamelCase              =						key.replace(f'''layer_norm{idx}'''    ,       f'''layer_norm.{int(lowerCAmelCase					)-1}'''					)
              if "layer_norm1" in key:
                     _UpperCamelCase              =						key.replace('''layer_norm1'''    ,       '''layer_norm_1'''					)
              if "layer_norm2" in key:
                     _UpperCamelCase              =						key.replace('''layer_norm2'''    ,       '''layer_norm_2'''					)
              if "block" in key:
                     # replace for example block1 by block.0
                     _UpperCamelCase              =						key[key.find('''block'''					) + len('''block'''					)]
                     _UpperCamelCase              =						key.replace(f'''block{idx}'''    ,       f'''block.{int(lowerCAmelCase					)-1}'''					)
              if "attn.q" in key:
                     _UpperCamelCase              =						key.replace('''attn.q'''    ,       '''attention.self.query'''					)
              if "attn.proj" in key:
                     _UpperCamelCase              =						key.replace('''attn.proj'''    ,       '''attention.output.dense'''					)
              if "attn" in key:
                     _UpperCamelCase              =						key.replace('''attn'''    ,       '''attention.self'''					)
              if "fc1" in key:
                     _UpperCamelCase              =						key.replace('''fc1'''    ,       '''dense1'''					)
              if "fc2" in key:
                     _UpperCamelCase              =						key.replace('''fc2'''    ,       '''dense2'''					)
              if "linear_pred" in key:
                     _UpperCamelCase              =						key.replace('''linear_pred'''    ,       '''classifier'''					)
              if "linear_fuse" in key:
                     _UpperCamelCase              =						key.replace('''linear_fuse.conv'''    ,       '''linear_fuse'''					)
                     _UpperCamelCase              =						key.replace('''linear_fuse.bn'''    ,       '''batch_norm'''					)
              if "linear_c" in key:
                     # replace for example linear_c4 by linear_c.3
                     _UpperCamelCase              =						key[key.find('''linear_c'''					) + len('''linear_c'''					)]
                     _UpperCamelCase              =						key.replace(f'''linear_c{idx}'''    ,       f'''linear_c.{int(lowerCAmelCase					)-1}'''					)
              if key.startswith('''head'''					):
                     _UpperCamelCase              =						key.replace('''head'''    ,       '''classifier'''					)
              _UpperCamelCase              =						value

       return new_state_dict


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase					):
       # for each of the encoder blocks:
       for i in range(config.num_encoder_blocks					):
              for j in range(config.depths[i]					):
                     # read in weights + bias of keys and values (which is a single matrix in the original implementation)
                     _UpperCamelCase              =						state_dict.pop(f'''segformer.encoder.block.{i}.{j}.attention.self.kv.weight'''					)
                     _UpperCamelCase              =						state_dict.pop(f'''segformer.encoder.block.{i}.{j}.attention.self.kv.bias'''					)
                     # next, add keys and values (in that order) to the state dict
                     _UpperCamelCase              =						kv_weight[
                         : config.hidden_sizes[i], :
                     ]
                     _UpperCamelCase              =						kv_bias[: config.hidden_sizes[i]]
                     _UpperCamelCase              =						kv_weight[
                         config.hidden_sizes[i] :, :
                     ]
                     _UpperCamelCase              =						kv_bias[
                         config.hidden_sizes[i] :
                     ]


def     SCREAMING_SNAKE_CASE			(      ):
       _UpperCamelCase              =						'''http://images.cocodataset.org/val2017/000000039769.jpg'''
       _UpperCamelCase              =						Image.open(requests.get(lowerCAmelCase    ,       stream=lowerCAmelCase					).raw					)

       return image


@torch.no_grad()
def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase    ,       lowerCAmelCase    ,       lowerCAmelCase					):
       _UpperCamelCase              =						SegformerConfig()
       _UpperCamelCase              =						False

       # set attributes based on model_name
       _UpperCamelCase              =						'''huggingface/label-files'''
       if "segformer" in model_name:
              _UpperCamelCase              =						model_name[len('''segformer.'''					) : len('''segformer.'''					) + 2]
              if "ade" in model_name:
                     _UpperCamelCase              =						150
                     _UpperCamelCase              =						'''ade20k-id2label.json'''
                     _UpperCamelCase              =						(1, 150, 128, 128)
              elif "city" in model_name:
                     _UpperCamelCase              =						19
                     _UpperCamelCase              =						'''cityscapes-id2label.json'''
                     _UpperCamelCase              =						(1, 19, 128, 128)
              else:
                     raise ValueError(f'''Model {model_name} not supported'''					)
       elif "mit" in model_name:
              _UpperCamelCase              =						True
              _UpperCamelCase              =						model_name[4:6]
              _UpperCamelCase              =						1_000
              _UpperCamelCase              =						'''imagenet-1k-id2label.json'''
              _UpperCamelCase              =						(1, 1_000)
       else:
              raise ValueError(f'''Model {model_name} not supported'''					)

       # set config attributes
       _UpperCamelCase              =						json.load(open(hf_hub_download(lowerCAmelCase    ,       lowerCAmelCase    ,       repo_type='''dataset'''					)    ,       '''r'''					)					)
       _UpperCamelCase              =						{int(lowerCAmelCase					): v for k, v in idalabel.items()}
       _UpperCamelCase              =						idalabel
       _UpperCamelCase              =						{v: k for k, v in idalabel.items()}
       if size == "b0":
              pass
       elif size == "b1":
              _UpperCamelCase              =						[64, 128, 320, 512]
              _UpperCamelCase              =						256
       elif size == "b2":
              _UpperCamelCase              =						[64, 128, 320, 512]
              _UpperCamelCase              =						768
              _UpperCamelCase              =						[3, 4, 6, 3]
       elif size == "b3":
              _UpperCamelCase              =						[64, 128, 320, 512]
              _UpperCamelCase              =						768
              _UpperCamelCase              =						[3, 4, 18, 3]
       elif size == "b4":
              _UpperCamelCase              =						[64, 128, 320, 512]
              _UpperCamelCase              =						768
              _UpperCamelCase              =						[3, 8, 27, 3]
       elif size == "b5":
              _UpperCamelCase              =						[64, 128, 320, 512]
              _UpperCamelCase              =						768
              _UpperCamelCase              =						[3, 6, 40, 3]
       else:
              raise ValueError(f'''Size {size} not supported'''					)

       # load image processor (only resize + normalize)
       _UpperCamelCase              =						SegformerImageProcessor(
           image_scale=(512, 512)    ,       keep_ratio=lowerCAmelCase    ,       align=lowerCAmelCase    ,       do_random_crop=lowerCAmelCase					)

       # prepare image
       _UpperCamelCase              =						prepare_img()
       _UpperCamelCase              =						image_processor(images=lowerCAmelCase    ,       return_tensors='''pt'''					).pixel_values

       logger.info(f'''Converting model {model_name}...'''					)

       # load original state dict
       if encoder_only:
              _UpperCamelCase              =						torch.load(lowerCAmelCase    ,       map_location=torch.device('''cpu'''					)					)
       else:
              _UpperCamelCase              =						torch.load(lowerCAmelCase    ,       map_location=torch.device('''cpu'''					)					)['''state_dict''']

       # rename keys
       _UpperCamelCase              =						rename_keys(lowerCAmelCase    ,       encoder_only=lowerCAmelCase					)
       if not encoder_only:
              del state_dict["decode_head.conv_seg.weight"]
              del state_dict["decode_head.conv_seg.bias"]

       # key and value matrices need special treatment
       read_in_k_v(lowerCAmelCase    ,       lowerCAmelCase					)

       # create HuggingFace model and load state dict
       if encoder_only:
              _UpperCamelCase              =						False
              _UpperCamelCase              =						SegformerForImageClassification(lowerCAmelCase					)
       else:
              _UpperCamelCase              =						SegformerForSemanticSegmentation(lowerCAmelCase					)
       model.load_state_dict(lowerCAmelCase					)
       model.eval()

       # forward pass
       _UpperCamelCase              =						model(lowerCAmelCase					)
       _UpperCamelCase              =						outputs.logits

       # set expected_slice based on model name
       # ADE20k checkpoints
       if model_name == "segformer.b0.512x512.ade.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-4.6_3_1_0, -5.5_2_3_2, -6.2_3_5_6], [-5.1_9_2_1, -6.1_4_4_4, -6.5_9_9_6], [-5.4_4_2_4, -6.2_7_9_0, -6.7_5_7_4]],
                      [[-12.1_391, -13.3_122, -13.9_554], [-12.8_732, -13.9_352, -14.3_563], [-12.9_438, -13.8_226, -14.2_513]],
                      [[-12.5_134, -13.4_686, -14.4_915], [-12.8_669, -14.4_343, -14.7_758], [-13.2_523, -14.5_819, -15.0_694]],
                  ]					)
       elif model_name == "segformer.b1.512x512.ade.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-7.5_8_2_0, -8.7_2_3_1, -8.3_2_1_5], [-8.0_6_0_0, -10.3_529, -10.0_304], [-7.5_2_0_8, -9.4_1_0_3, -9.6_2_3_9]],
                      [[-12.6_918, -13.8_994, -13.7_137], [-13.3_196, -15.7_523, -15.4_789], [-12.9_343, -14.8_757, -14.9_689]],
                      [[-11.1_911, -11.9_421, -11.3_243], [-11.3_342, -13.6_839, -13.3_581], [-10.3_909, -12.1_832, -12.4_858]],
                  ]					)
       elif model_name == "segformer.b2.512x512.ade.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-11.8_173, -14.3_850, -16.3_128], [-14.5_648, -16.5_804, -18.6_568], [-14.7_223, -15.7_387, -18.4_218]],
                      [[-15.7_290, -17.9_171, -19.4_423], [-18.3_105, -19.9_448, -21.4_661], [-17.9_296, -18.6_497, -20.7_910]],
                      [[-15.0_783, -17.0_336, -18.2_789], [-16.8_771, -18.6_870, -20.1_612], [-16.2_454, -17.1_426, -19.5_055]],
                  ]					)
       elif model_name == "segformer.b3.512x512.ade.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-9.0_8_7_8, -10.2_081, -10.1_891], [-9.3_1_4_4, -10.7_941, -10.9_843], [-9.2_2_9_4, -10.3_855, -10.5_704]],
                      [[-12.2_316, -13.9_068, -13.6_102], [-12.9_161, -14.3_702, -14.3_235], [-12.5_233, -13.7_174, -13.7_932]],
                      [[-14.6_275, -15.2_490, -14.9_727], [-14.3_400, -15.9_687, -16.2_827], [-14.1_484, -15.4_033, -15.8_937]],
                  ]					)
       elif model_name == "segformer.b4.512x512.ade.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-12.3_144, -13.2_447, -14.0_802], [-13.3_614, -14.5_816, -15.6_117], [-13.3_340, -14.4_433, -16.2_219]],
                      [[-19.2_781, -20.4_128, -20.7_506], [-20.6_153, -21.6_566, -22.0_998], [-19.9_800, -21.0_430, -22.1_494]],
                      [[-18.8_739, -19.7_804, -21.1_834], [-20.1_233, -21.6_765, -23.2_944], [-20.0_315, -21.2_641, -23.6_944]],
                  ]					)
       elif model_name == "segformer.b5.640x640.ade.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-9.5_5_2_4, -12.0_835, -11.7_348], [-10.5_229, -13.6_446, -14.5_662], [-9.5_8_4_2, -12.8_851, -13.9_414]],
                      [[-15.3_432, -17.5_323, -17.0_818], [-16.3_330, -18.9_255, -19.2_101], [-15.1_340, -17.7_848, -18.3_971]],
                      [[-12.6_072, -14.9_486, -14.6_631], [-13.7_629, -17.0_907, -17.7_745], [-12.7_899, -16.1_695, -17.1_671]],
                  ]					)
       # Cityscapes checkpoints
       elif model_name == "segformer.b0.1024x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-11.9_295, -13.4_057, -14.8_106], [-13.3_431, -14.8_179, -15.3_781], [-14.2_836, -15.5_942, -16.1_588]],
                      [[-11.4_906, -12.8_067, -13.6_564], [-13.1_189, -14.0_500, -14.1_543], [-13.8_748, -14.5_136, -14.8_789]],
                      [[0.5_3_7_4, 0.1_0_6_7, -0.4_7_4_2], [0.1_1_4_1, -0.2_2_5_5, -0.7_0_9_9], [-0.3_0_0_0, -0.5_9_2_4, -1.3_1_0_5]],
                  ]					)
       elif model_name == "segformer.b0.512x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-7.8_2_1_7, -9.8_7_6_7, -10.1_717], [-9.4_4_3_8, -10.9_058, -11.4_047], [-9.7_9_3_9, -12.3_495, -12.1_079]],
                      [[-7.1_5_1_4, -9.5_3_3_6, -10.0_860], [-9.7_7_7_6, -11.6_822, -11.8_439], [-10.1_411, -12.7_655, -12.8_972]],
                      [[0.3_0_2_1, 0.0_8_0_5, -0.2_3_1_0], [-0.0_3_2_8, -0.1_6_0_5, -0.2_7_1_4], [-0.1_4_0_8, -0.5_4_7_7, -0.6_9_7_6]],
                  ]					)
       elif model_name == "segformer.b0.640x1280.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [
                          [-1.1_372E01, -1.2_787E01, -1.3_477E01],
                          [-1.2_536E01, -1.4_194E01, -1.4_409E01],
                          [-1.3_217E01, -1.4_888E01, -1.5_327E01],
                      ],
                      [
                          [-1.4_791E01, -1.7_122E01, -1.8_277E01],
                          [-1.7_163E01, -1.9_192E01, -1.9_533E01],
                          [-1.7_897E01, -1.9_991E01, -2.0_315E01],
                      ],
                      [
                          [7.6_723E-01, 4.1_921E-01, -7.7_878E-02],
                          [4.7_772E-01, 9.5_557E-03, -2.8_082E-01],
                          [3.6_032E-01, -2.4_826E-01, -5.1_168E-01],
                      ],
                  ]					)
       elif model_name == "segformer.b0.768x768.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-9.4_9_5_9, -11.3_087, -11.7_479], [-11.0_025, -12.6_540, -12.3_319], [-11.4_064, -13.0_487, -12.9_905]],
                      [[-9.8_9_0_5, -11.3_084, -12.0_854], [-11.1_726, -12.7_698, -12.9_583], [-11.5_985, -13.3_278, -14.1_774]],
                      [[0.2_2_1_3, 0.0_1_9_2, -0.2_4_6_6], [-0.1_7_3_1, -0.4_2_1_3, -0.4_8_7_4], [-0.3_1_2_6, -0.6_5_4_1, -1.1_3_8_9]],
                  ]					)
       elif model_name == "segformer.b1.1024x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-13.5_748, -13.9_111, -12.6_500], [-14.3_500, -15.3_683, -14.2_328], [-14.7_532, -16.0_424, -15.6_087]],
                      [[-17.1_651, -15.8_725, -12.9_653], [-17.2_580, -17.3_718, -14.8_223], [-16.6_058, -16.8_783, -16.7_452]],
                      [[-3.6_4_5_6, -3.0_2_0_9, -1.4_2_0_3], [-3.0_7_9_7, -3.1_9_5_9, -2.0_0_0_0], [-1.8_7_5_7, -1.9_2_1_7, -1.6_9_9_7]],
                  ]					)
       elif model_name == "segformer.b2.1024x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-16.0_976, -16.4_856, -17.3_962], [-16.6_234, -19.0_342, -19.7_685], [-16.0_900, -18.0_661, -19.1_180]],
                      [[-18.4_750, -18.8_488, -19.5_074], [-19.4_030, -22.1_570, -22.5_977], [-19.1_191, -20.8_486, -22.3_783]],
                      [[-4.5_1_7_8, -5.5_0_3_7, -6.5_1_0_9], [-5.0_8_8_4, -7.2_1_7_4, -8.0_3_3_4], [-4.4_1_5_6, -5.8_1_1_7, -7.2_9_7_0]],
                  ]					)
       elif model_name == "segformer.b3.1024x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-14.2_081, -14.4_732, -14.1_977], [-14.5_867, -16.4_423, -16.6_356], [-13.4_441, -14.9_685, -16.8_696]],
                      [[-14.4_576, -14.7_073, -15.0_451], [-15.0_816, -17.6_237, -17.9_873], [-14.4_213, -16.0_199, -18.5_992]],
                      [[-4.7_3_4_9, -4.9_5_8_8, -5.0_9_6_6], [-4.3_2_1_0, -6.9_3_2_5, -7.2_5_9_1], [-3.4_3_1_2, -4.7_4_8_4, -7.1_9_1_7]],
                  ]					)
       elif model_name == "segformer.b4.1024x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-11.7_737, -11.9_526, -11.3_273], [-13.6_692, -14.4_574, -13.8_878], [-13.8_937, -14.6_924, -15.9_345]],
                      [[-14.6_706, -14.5_330, -14.1_306], [-16.1_502, -16.8_180, -16.4_269], [-16.8_338, -17.8_939, -20.1_746]],
                      [[1.0_4_9_1, 0.8_2_8_9, 1.0_3_1_0], [1.1_0_4_4, 0.5_2_1_9, 0.8_0_5_5], [1.0_8_9_9, 0.6_9_2_6, 0.5_5_9_0]],
                  ]					)
       elif model_name == "segformer.b5.1024x1024.city.160k":
              _UpperCamelCase              =						torch.tensor(
                  [
                      [[-12.5_641, -13.4_777, -13.0_684], [-13.9_587, -15.8_983, -16.6_557], [-13.3_109, -15.7_350, -16.3_141]],
                      [[-14.7_074, -15.4_352, -14.5_944], [-16.6_353, -18.1_663, -18.6_120], [-15.1_702, -18.0_329, -18.1_547]],
                      [[-1.7_9_9_0, -2.0_9_5_1, -1.7_7_8_4], [-2.6_3_9_7, -3.8_2_4_5, -3.9_6_8_6], [-1.5_2_6_4, -2.8_1_2_6, -2.9_3_1_6]],
                  ]					)
       else:
              _UpperCamelCase              =						logits.argmax(-1					).item()
              print('''Predicted class:'''    ,       model.config.idalabel[predicted_class_idx]					)

       # verify logits
       if not encoder_only:
              assert logits.shape == expected_shape
              assert torch.allclose(logits[0, :3, :3, :3]    ,       lowerCAmelCase    ,       atol=1E-2					)

       # finally, save model and image processor
       logger.info(f'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'''					)
       Path(lowerCAmelCase					).mkdir(exist_ok=lowerCAmelCase					)
       model.save_pretrained(lowerCAmelCase					)
       image_processor.save_pretrained(lowerCAmelCase					)


if __name__ == "__main__":
       lowercase    : Any    	=  argparse.ArgumentParser()

       parser.add_argument(
           """--model_name""",
           default="""segformer.b0.512x512.ade.160k""",
           type=str,
           help="""Name of the model you'd like to convert.""",
       )
       parser.add_argument(
           """--checkpoint_path""", default=None, type=str, help="""Path to the original PyTorch checkpoint (.pth file)."""
       )
       parser.add_argument(
           """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
       )
       lowercase    : str    	=  parser.parse_args()
       convert_segformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)













from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch

@require_torch
@require_tf
class  __A(						__UpperCAmelCase ):

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						SMALL_MODEL_IDENTIFIER
            _UpperCamelCase              =						'''pt'''
            _UpperCamelCase              =						'''tf'''

     def 	_UpperCamelCase     (      self,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						AutoModel.from_pretrained(self.test_model      )
            model_pt.save_pretrained(A      )

     def 	_UpperCamelCase     (      self,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						TFAutoModel.from_pretrained(self.test_model,       from_pt=A      )
            model_tf.save_pretrained(A      )

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						'''mock_framework'''

            # Framework provided - return whatever the user provides
            _UpperCamelCase              =						FeaturesManager.determine_framework(self.test_model,       A      )
            self.assertEqual(A,       A      )

            # Local checkpoint and framework provided - return provided framework
            # PyTorch checkpoint
            with TemporaryDirectory() as local_pt_ckpt:
                   self._setup_pt_ckpt(A      )
                   _UpperCamelCase              =						FeaturesManager.determine_framework(A,       A      )
                   self.assertEqual(A,       A      )

            # TensorFlow checkpoint
            with TemporaryDirectory() as local_tf_ckpt:
                   self._setup_tf_ckpt(A      )
                   _UpperCamelCase              =						FeaturesManager.determine_framework(A,       A      )
                   self.assertEqual(A,       A      )

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            with TemporaryDirectory() as local_pt_ckpt:
                   self._setup_pt_ckpt(A      )
                   _UpperCamelCase              =						FeaturesManager.determine_framework(A      )
                   self.assertEqual(A,       self.framework_pt      )

            # TensorFlow checkpoint
            with TemporaryDirectory() as local_tf_ckpt:
                   self._setup_tf_ckpt(A      )
                   _UpperCamelCase              =						FeaturesManager.determine_framework(A      )
                   self.assertEqual(A,       self.framework_tf      )

            # Invalid local checkpoint
            with TemporaryDirectory() as local_invalid_ckpt:
                   with self.assertRaises(A      ):
                          _UpperCamelCase              =						FeaturesManager.determine_framework(A      )



     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						MagicMock(return_value=A      )
            with patch('''transformers.onnx.features.is_tf_available''',       A      ):
                   _UpperCamelCase              =						FeaturesManager.determine_framework(self.test_model      )
                   self.assertEqual(A,       self.framework_pt      )

            # PyTorch not in environment -> use TensorFlow
            _UpperCamelCase              =						MagicMock(return_value=A      )
            with patch('''transformers.onnx.features.is_torch_available''',       A      ):
                   _UpperCamelCase              =						FeaturesManager.determine_framework(self.test_model      )
                   self.assertEqual(A,       self.framework_tf      )

            # Both in environment -> use PyTorch
            _UpperCamelCase              =						MagicMock(return_value=A      )
            _UpperCamelCase              =						MagicMock(return_value=A      )
            with patch('''transformers.onnx.features.is_tf_available''',       A      ), patch(
                '''transformers.onnx.features.is_torch_available''',       A      ):
                   _UpperCamelCase              =						FeaturesManager.determine_framework(self.test_model      )
                   self.assertEqual(A,       self.framework_pt      )

            # Both not in environment -> raise error
            _UpperCamelCase              =						MagicMock(return_value=A      )
            _UpperCamelCase              =						MagicMock(return_value=A      )
            with patch('''transformers.onnx.features.is_tf_available''',       A      ), patch(
                '''transformers.onnx.features.is_torch_available''',       A      ):
                   with self.assertRaises(A      ):
                          _UpperCamelCase              =						FeaturesManager.determine_framework(self.test_model      )






