





'''simple docstring'''






import argparse
import json
import pickle
from pathlib import Path

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

from transformers import MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
A__    :       List[str]              =logging.get_logger(__name__)
def        UpperCamelCase__				(		lowerCAmelCase  ):





      """simple docstring"""


      _lowerCAmelCase												=							SwinConfig.from_pretrained(
          """microsoft/swin-tiny-patch4-window7-224"""							,     out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]  )
      _lowerCAmelCase												=							MaskFormerConfig(backbone_config=lowerCAmelCase  )

      _lowerCAmelCase												=							"""huggingface/label-files"""
      if "ade20k-full" in model_name:
            # this should be ok
            _lowerCAmelCase												=							8_47
            _lowerCAmelCase												=							"""maskformer-ade20k-full-id2label.json"""
      elif "ade" in model_name:
            # this should be ok
            _lowerCAmelCase												=							1_50
            _lowerCAmelCase												=							"""ade20k-id2label.json"""
      elif "coco-stuff" in model_name:
            # this should be ok
            _lowerCAmelCase												=							1_71
            _lowerCAmelCase												=							"""maskformer-coco-stuff-id2label.json"""
      elif "coco" in model_name:
            # TODO
            _lowerCAmelCase												=							1_33
            _lowerCAmelCase												=							"""coco-panoptic-id2label.json"""
      elif "cityscapes" in model_name:
            # this should be ok
            _lowerCAmelCase												=							19
            _lowerCAmelCase												=							"""cityscapes-id2label.json"""
      elif "vistas" in model_name:
            # this should be ok
            _lowerCAmelCase												=							65
            _lowerCAmelCase												=							"""mapillary-vistas-id2label.json"""

      _lowerCAmelCase												=							json.load(open(hf_hub_download(lowerCAmelCase							,     lowerCAmelCase							,     repo_type="""dataset"""  )							,     """r"""  )  )
      _lowerCAmelCase												=							{int(lowerCAmelCase  ): v for k, v in idalabel.items()}

      return config
def        UpperCamelCase__				(		lowerCAmelCase  ):





      """simple docstring"""


      _lowerCAmelCase												=							[]
      # stem
      # fmt: off
      rename_keys.append(("""backbone.patch_embed.proj.weight""", """model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.weight""")  )
      rename_keys.append(("""backbone.patch_embed.proj.bias""", """model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.bias""")  )
      rename_keys.append(("""backbone.patch_embed.norm.weight""", """model.pixel_level_module.encoder.model.embeddings.norm.weight""")  )
      rename_keys.append(("""backbone.patch_embed.norm.bias""", """model.pixel_level_module.encoder.model.embeddings.norm.bias""")  )
      # stages
      for i in range(len(config.backbone_config.depths  )  ):
            for j in range(config.backbone_config.depths[i]  ):
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm1.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm1.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.relative_position_bias_table", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.relative_position_index", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.proj.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.attn.proj.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm2.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.norm2.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc1.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc1.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc2.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.blocks.{j}.mlp.fc2.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.bias")  )

            if i < 3:
                  rename_keys.append((f"backbone.layers.{i}.downsample.reduction.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.reduction.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.downsample.norm.weight", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.weight")  )
                  rename_keys.append((f"backbone.layers.{i}.downsample.norm.bias", f"model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.bias")  )
            rename_keys.append((f"backbone.norm{i}.weight", f"model.pixel_level_module.encoder.hidden_states_norms.{i}.weight")  )
            rename_keys.append((f"backbone.norm{i}.bias", f"model.pixel_level_module.encoder.hidden_states_norms.{i}.bias")  )

      # FPN
      rename_keys.append(("""sem_seg_head.layer_4.weight""", """model.pixel_level_module.decoder.fpn.stem.0.weight""")  )
      rename_keys.append(("""sem_seg_head.layer_4.norm.weight""", """model.pixel_level_module.decoder.fpn.stem.1.weight""")  )
      rename_keys.append(("""sem_seg_head.layer_4.norm.bias""", """model.pixel_level_module.decoder.fpn.stem.1.bias""")  )
      for source_index, target_index in zip(range(3							,     0							,     -1  )							,     range(0							,     3  )  ):
            rename_keys.append((f"sem_seg_head.adapter_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight")  )
            rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight")  )
            rename_keys.append((f"sem_seg_head.adapter_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias")  )
            rename_keys.append((f"sem_seg_head.layer_{source_index}.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight")  )
            rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.weight", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight")  )
            rename_keys.append((f"sem_seg_head.layer_{source_index}.norm.bias", f"model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias")  )
      rename_keys.append(("""sem_seg_head.mask_features.weight""", """model.pixel_level_module.decoder.mask_projection.weight""")  )
      rename_keys.append(("""sem_seg_head.mask_features.bias""", """model.pixel_level_module.decoder.mask_projection.bias""")  )

      # Transformer decoder
      for idx in range(config.decoder_config.decoder_layers  ):
            # self-attention out projection
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias")  )
            # cross-attention out projection
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias")  )
            # MLP 1
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight", f"model.transformer_module.decoder.layers.{idx}.fc1.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias", f"model.transformer_module.decoder.layers.{idx}.fc1.bias")  )
            # MLP 2
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight", f"model.transformer_module.decoder.layers.{idx}.fc2.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias", f"model.transformer_module.decoder.layers.{idx}.fc2.bias")  )
            # layernorm 1 (self-attention layernorm)
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias", f"model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias")  )
            # layernorm 2 (cross-attention layernorm)
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias", f"model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias")  )
            # layernorm 3 (final layernorm)
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias", f"model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias")  )

      rename_keys.append(("""sem_seg_head.predictor.transformer.decoder.norm.weight""", """model.transformer_module.decoder.layernorm.weight""")  )
      rename_keys.append(("""sem_seg_head.predictor.transformer.decoder.norm.bias""", """model.transformer_module.decoder.layernorm.bias""")  )

      # heads on top
      rename_keys.append(("""sem_seg_head.predictor.query_embed.weight""", """model.transformer_module.queries_embedder.weight""")  )

      rename_keys.append(("""sem_seg_head.predictor.input_proj.weight""", """model.transformer_module.input_projection.weight""")  )
      rename_keys.append(("""sem_seg_head.predictor.input_proj.bias""", """model.transformer_module.input_projection.bias""")  )

      rename_keys.append(("""sem_seg_head.predictor.class_embed.weight""", """class_predictor.weight""")  )
      rename_keys.append(("""sem_seg_head.predictor.class_embed.bias""", """class_predictor.bias""")  )

      for i in range(3  ):
            rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.weight", f"mask_embedder.{i}.0.weight")  )
            rename_keys.append((f"sem_seg_head.predictor.mask_embed.layers.{i}.bias", f"mask_embedder.{i}.0.bias")  )
      # fmt: on

      return rename_keys
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      _lowerCAmelCase												=							dct.pop(lowerCAmelCase  )
      _lowerCAmelCase												=							val
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      _lowerCAmelCase												=							[int(backbone_config.embed_dim * 2**i  ) for i in range(len(backbone_config.depths  )  )]
      for i in range(len(backbone_config.depths  )  ):
            _lowerCAmelCase												=							num_features[i]
            for j in range(backbone_config.depths[i]  ):
                  # fmt: off
                  # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
                  _lowerCAmelCase												=							state_dict.pop(f"backbone.layers.{i}.blocks.{j}.attn.qkv.weight"  )
                  _lowerCAmelCase												=							state_dict.pop(f"backbone.layers.{i}.blocks.{j}.attn.qkv.bias"  )
                  # next, add query, keys and values (in that order) to the state dict
                  _lowerCAmelCase												=							in_proj_weight[:dim, :]
                  _lowerCAmelCase												=							in_proj_bias[: dim]
                  _lowerCAmelCase												=							in_proj_weight[
                      dim : dim * 2, :
                  ]
                  _lowerCAmelCase												=							in_proj_bias[
                      dim : dim * 2
                  ]
                  _lowerCAmelCase												=							in_proj_weight[
                      -dim :, :
                  ]
                  _lowerCAmelCase												=							in_proj_bias[-dim :]
                  # fmt: on
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      _lowerCAmelCase												=							config.decoder_config.hidden_size
      for idx in range(config.decoder_config.decoder_layers  ):
            # read in weights + bias of self-attention input projection layer (in the original implementation, this is a single matrix + bias)
            _lowerCAmelCase												=							state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight"  )
            _lowerCAmelCase												=							state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_bias"  )
            # next, add query, keys and values (in that order) to the state dict
            _lowerCAmelCase												=							in_proj_weight[: hidden_size, :]
            _lowerCAmelCase												=							in_proj_bias[:config.hidden_size]
            _lowerCAmelCase												=							in_proj_weight[hidden_size : hidden_size * 2, :]
            _lowerCAmelCase												=							in_proj_bias[hidden_size : hidden_size * 2]
            _lowerCAmelCase												=							in_proj_weight[-hidden_size :, :]
            _lowerCAmelCase												=							in_proj_bias[-hidden_size :]
            # read in weights + bias of cross-attention input projection layer (in the original implementation, this is a single matrix + bias)
            _lowerCAmelCase												=							state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight"  )
            _lowerCAmelCase												=							state_dict.pop(f"sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_bias"  )
            # next, add query, keys and values (in that order) to the state dict
            _lowerCAmelCase												=							in_proj_weight[: hidden_size, :]
            _lowerCAmelCase												=							in_proj_bias[:config.hidden_size]
            _lowerCAmelCase												=							in_proj_weight[hidden_size : hidden_size * 2, :]
            _lowerCAmelCase												=							in_proj_bias[hidden_size : hidden_size * 2]
            _lowerCAmelCase												=							in_proj_weight[-hidden_size :, :]
            _lowerCAmelCase												=							in_proj_bias[-hidden_size :]
      # fmt: on
def        UpperCamelCase__				(		):





      """simple docstring"""


      _lowerCAmelCase												=							"""http://images.cocodataset.org/val2017/000000039769.jpg"""
      _lowerCAmelCase												=							Image.open(requests.get(lowerCAmelCase							,     stream=lowerCAmelCase  ).raw  )
      return im
@torch.no_grad()
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase = False  ):





      """simple docstring"""


      _lowerCAmelCase												=							get_maskformer_config(lowerCAmelCase  )

      # load original state_dict
      with open(lowerCAmelCase							,     """rb"""  ) as f:
            _lowerCAmelCase												=							pickle.load(lowerCAmelCase  )
      _lowerCAmelCase												=							data["""model"""]

      # for name, param in state_dict.items():
      #     print(name, param.shape)

      # rename keys
      _lowerCAmelCase												=							create_rename_keys(lowerCAmelCase  )
      for src, dest in rename_keys:
            rename_key(lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  )
      read_in_swin_q_k_v(lowerCAmelCase							,     config.backbone_config  )
      read_in_decoder_q_k_v(lowerCAmelCase							,     lowerCAmelCase  )

      # update to torch tensors
      for key, value in state_dict.items():
            _lowerCAmelCase												=							torch.from_numpy(lowerCAmelCase  )

      # load 🤗 model
      _lowerCAmelCase												=							MaskFormerForInstanceSegmentation(lowerCAmelCase  )
      model.eval()

      for name, param in model.named_parameters():
            print(lowerCAmelCase							,     param.shape  )

      _lowerCAmelCase					,							_lowerCAmelCase												=							model.load_state_dict(lowerCAmelCase							,     strict=lowerCAmelCase  )
      assert missing_keys == [
          "model.pixel_level_module.encoder.model.layernorm.weight",
          "model.pixel_level_module.encoder.model.layernorm.bias",
      ]
      assert len(lowerCAmelCase  ) == 0, f"Unexpected keys: {unexpected_keys}"

      # verify results
      _lowerCAmelCase												=							prepare_img()
      if "vistas" in model_name:
            _lowerCAmelCase												=							65
      elif "cityscapes" in model_name:
            _lowerCAmelCase												=							6_55_35
      else:
            _lowerCAmelCase												=							2_55
      _lowerCAmelCase												=							True if """ade""" in model_name else False
      _lowerCAmelCase												=							MaskFormerImageProcessor(ignore_index=lowerCAmelCase							,     reduce_labels=lowerCAmelCase  )

      _lowerCAmelCase												=							image_processor(lowerCAmelCase							,     return_tensors="""pt"""  )

      _lowerCAmelCase												=							model(**lowerCAmelCase  )

      print("""Logits:"""							,     outputs.class_queries_logits[0, :3, :3]  )

      if model_name == "maskformer-swin-tiny-ade":
            _lowerCAmelCase												=							torch.tensor(
                [[3.6_353, -4.4_770, -2.6_065], [0.5_081, -4.2_394, -3.5_343], [2.1_909, -5.0_353, -1.9_323]]  )
      assert torch.allclose(outputs.class_queries_logits[0, :3, :3]							,     lowerCAmelCase							,     atol=1e-4  )
      print("""Looks ok!"""  )

      if pytorch_dump_folder_path is not None:
            print(f"Saving 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 push_to_hub:
            print("""Pushing model and image processor to the hub..."""  )
            model.push_to_hub(f"nielsr/{model_name}"  )
            image_processor.push_to_hub(f"nielsr/{model_name}"  )


if __name__ == "__main__":
      A__    :       str              =argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '''--model_name''',
          default='''maskformer-swin-tiny-ade''',
          type=str,
          help=('''Name of the MaskFormer model you\'d like to convert''',),
      )
      parser.add_argument(
          '''--checkpoint_path''',
          default='''/Users/nielsrogge/Documents/MaskFormer_checkpoints/MaskFormer-Swin-tiny-ADE20k/model.pkl''',
          type=str,
          help='''Path to the original state dict (.pth file).''',
      )
      parser.add_argument(
          '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
      )
      parser.add_argument(
          '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
      )

      A__    :       List[Any]              =parser.parse_args()
      convert_maskformer_checkpoint(
          args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
      )













'''simple docstring'''






# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.auto import AutoModelForSeqaSeqLM, AutoTokenizer
from .base import PipelineTool


A__    :       List[str]              ={
    '''Acehnese Arabic''': '''ace_Arab''',
    '''Acehnese Latin''': '''ace_Latn''',
    '''Mesopotamian Arabic''': '''acm_Arab''',
    '''Ta\'izzi-Adeni Arabic''': '''acq_Arab''',
    '''Tunisian Arabic''': '''aeb_Arab''',
    '''Afrikaans''': '''afr_Latn''',
    '''South Levantine Arabic''': '''ajp_Arab''',
    '''Akan''': '''aka_Latn''',
    '''Amharic''': '''amh_Ethi''',
    '''North Levantine Arabic''': '''apc_Arab''',
    '''Modern Standard Arabic''': '''arb_Arab''',
    '''Modern Standard Arabic Romanized''': '''arb_Latn''',
    '''Najdi Arabic''': '''ars_Arab''',
    '''Moroccan Arabic''': '''ary_Arab''',
    '''Egyptian Arabic''': '''arz_Arab''',
    '''Assamese''': '''asm_Beng''',
    '''Asturian''': '''ast_Latn''',
    '''Awadhi''': '''awa_Deva''',
    '''Central Aymara''': '''ayr_Latn''',
    '''South Azerbaijani''': '''azb_Arab''',
    '''North Azerbaijani''': '''azj_Latn''',
    '''Bashkir''': '''bak_Cyrl''',
    '''Bambara''': '''bam_Latn''',
    '''Balinese''': '''ban_Latn''',
    '''Belarusian''': '''bel_Cyrl''',
    '''Bemba''': '''bem_Latn''',
    '''Bengali''': '''ben_Beng''',
    '''Bhojpuri''': '''bho_Deva''',
    '''Banjar Arabic''': '''bjn_Arab''',
    '''Banjar Latin''': '''bjn_Latn''',
    '''Standard Tibetan''': '''bod_Tibt''',
    '''Bosnian''': '''bos_Latn''',
    '''Buginese''': '''bug_Latn''',
    '''Bulgarian''': '''bul_Cyrl''',
    '''Catalan''': '''cat_Latn''',
    '''Cebuano''': '''ceb_Latn''',
    '''Czech''': '''ces_Latn''',
    '''Chokwe''': '''cjk_Latn''',
    '''Central Kurdish''': '''ckb_Arab''',
    '''Crimean Tatar''': '''crh_Latn''',
    '''Welsh''': '''cym_Latn''',
    '''Danish''': '''dan_Latn''',
    '''German''': '''deu_Latn''',
    '''Southwestern Dinka''': '''dik_Latn''',
    '''Dyula''': '''dyu_Latn''',
    '''Dzongkha''': '''dzo_Tibt''',
    '''Greek''': '''ell_Grek''',
    '''English''': '''eng_Latn''',
    '''Esperanto''': '''epo_Latn''',
    '''Estonian''': '''est_Latn''',
    '''Basque''': '''eus_Latn''',
    '''Ewe''': '''ewe_Latn''',
    '''Faroese''': '''fao_Latn''',
    '''Fijian''': '''fij_Latn''',
    '''Finnish''': '''fin_Latn''',
    '''Fon''': '''fon_Latn''',
    '''French''': '''fra_Latn''',
    '''Friulian''': '''fur_Latn''',
    '''Nigerian Fulfulde''': '''fuv_Latn''',
    '''Scottish Gaelic''': '''gla_Latn''',
    '''Irish''': '''gle_Latn''',
    '''Galician''': '''glg_Latn''',
    '''Guarani''': '''grn_Latn''',
    '''Gujarati''': '''guj_Gujr''',
    '''Haitian Creole''': '''hat_Latn''',
    '''Hausa''': '''hau_Latn''',
    '''Hebrew''': '''heb_Hebr''',
    '''Hindi''': '''hin_Deva''',
    '''Chhattisgarhi''': '''hne_Deva''',
    '''Croatian''': '''hrv_Latn''',
    '''Hungarian''': '''hun_Latn''',
    '''Armenian''': '''hye_Armn''',
    '''Igbo''': '''ibo_Latn''',
    '''Ilocano''': '''ilo_Latn''',
    '''Indonesian''': '''ind_Latn''',
    '''Icelandic''': '''isl_Latn''',
    '''Italian''': '''ita_Latn''',
    '''Javanese''': '''jav_Latn''',
    '''Japanese''': '''jpn_Jpan''',
    '''Kabyle''': '''kab_Latn''',
    '''Jingpho''': '''kac_Latn''',
    '''Kamba''': '''kam_Latn''',
    '''Kannada''': '''kan_Knda''',
    '''Kashmiri Arabic''': '''kas_Arab''',
    '''Kashmiri Devanagari''': '''kas_Deva''',
    '''Georgian''': '''kat_Geor''',
    '''Central Kanuri Arabic''': '''knc_Arab''',
    '''Central Kanuri Latin''': '''knc_Latn''',
    '''Kazakh''': '''kaz_Cyrl''',
    '''Kabiyè''': '''kbp_Latn''',
    '''Kabuverdianu''': '''kea_Latn''',
    '''Khmer''': '''khm_Khmr''',
    '''Kikuyu''': '''kik_Latn''',
    '''Kinyarwanda''': '''kin_Latn''',
    '''Kyrgyz''': '''kir_Cyrl''',
    '''Kimbundu''': '''kmb_Latn''',
    '''Northern Kurdish''': '''kmr_Latn''',
    '''Kikongo''': '''kon_Latn''',
    '''Korean''': '''kor_Hang''',
    '''Lao''': '''lao_Laoo''',
    '''Ligurian''': '''lij_Latn''',
    '''Limburgish''': '''lim_Latn''',
    '''Lingala''': '''lin_Latn''',
    '''Lithuanian''': '''lit_Latn''',
    '''Lombard''': '''lmo_Latn''',
    '''Latgalian''': '''ltg_Latn''',
    '''Luxembourgish''': '''ltz_Latn''',
    '''Luba-Kasai''': '''lua_Latn''',
    '''Ganda''': '''lug_Latn''',
    '''Luo''': '''luo_Latn''',
    '''Mizo''': '''lus_Latn''',
    '''Standard Latvian''': '''lvs_Latn''',
    '''Magahi''': '''mag_Deva''',
    '''Maithili''': '''mai_Deva''',
    '''Malayalam''': '''mal_Mlym''',
    '''Marathi''': '''mar_Deva''',
    '''Minangkabau Arabic ''': '''min_Arab''',
    '''Minangkabau Latin''': '''min_Latn''',
    '''Macedonian''': '''mkd_Cyrl''',
    '''Plateau Malagasy''': '''plt_Latn''',
    '''Maltese''': '''mlt_Latn''',
    '''Meitei Bengali''': '''mni_Beng''',
    '''Halh Mongolian''': '''khk_Cyrl''',
    '''Mossi''': '''mos_Latn''',
    '''Maori''': '''mri_Latn''',
    '''Burmese''': '''mya_Mymr''',
    '''Dutch''': '''nld_Latn''',
    '''Norwegian Nynorsk''': '''nno_Latn''',
    '''Norwegian Bokmål''': '''nob_Latn''',
    '''Nepali''': '''npi_Deva''',
    '''Northern Sotho''': '''nso_Latn''',
    '''Nuer''': '''nus_Latn''',
    '''Nyanja''': '''nya_Latn''',
    '''Occitan''': '''oci_Latn''',
    '''West Central Oromo''': '''gaz_Latn''',
    '''Odia''': '''ory_Orya''',
    '''Pangasinan''': '''pag_Latn''',
    '''Eastern Panjabi''': '''pan_Guru''',
    '''Papiamento''': '''pap_Latn''',
    '''Western Persian''': '''pes_Arab''',
    '''Polish''': '''pol_Latn''',
    '''Portuguese''': '''por_Latn''',
    '''Dari''': '''prs_Arab''',
    '''Southern Pashto''': '''pbt_Arab''',
    '''Ayacucho Quechua''': '''quy_Latn''',
    '''Romanian''': '''ron_Latn''',
    '''Rundi''': '''run_Latn''',
    '''Russian''': '''rus_Cyrl''',
    '''Sango''': '''sag_Latn''',
    '''Sanskrit''': '''san_Deva''',
    '''Santali''': '''sat_Olck''',
    '''Sicilian''': '''scn_Latn''',
    '''Shan''': '''shn_Mymr''',
    '''Sinhala''': '''sin_Sinh''',
    '''Slovak''': '''slk_Latn''',
    '''Slovenian''': '''slv_Latn''',
    '''Samoan''': '''smo_Latn''',
    '''Shona''': '''sna_Latn''',
    '''Sindhi''': '''snd_Arab''',
    '''Somali''': '''som_Latn''',
    '''Southern Sotho''': '''sot_Latn''',
    '''Spanish''': '''spa_Latn''',
    '''Tosk Albanian''': '''als_Latn''',
    '''Sardinian''': '''srd_Latn''',
    '''Serbian''': '''srp_Cyrl''',
    '''Swati''': '''ssw_Latn''',
    '''Sundanese''': '''sun_Latn''',
    '''Swedish''': '''swe_Latn''',
    '''Swahili''': '''swh_Latn''',
    '''Silesian''': '''szl_Latn''',
    '''Tamil''': '''tam_Taml''',
    '''Tatar''': '''tat_Cyrl''',
    '''Telugu''': '''tel_Telu''',
    '''Tajik''': '''tgk_Cyrl''',
    '''Tagalog''': '''tgl_Latn''',
    '''Thai''': '''tha_Thai''',
    '''Tigrinya''': '''tir_Ethi''',
    '''Tamasheq Latin''': '''taq_Latn''',
    '''Tamasheq Tifinagh''': '''taq_Tfng''',
    '''Tok Pisin''': '''tpi_Latn''',
    '''Tswana''': '''tsn_Latn''',
    '''Tsonga''': '''tso_Latn''',
    '''Turkmen''': '''tuk_Latn''',
    '''Tumbuka''': '''tum_Latn''',
    '''Turkish''': '''tur_Latn''',
    '''Twi''': '''twi_Latn''',
    '''Central Atlas Tamazight''': '''tzm_Tfng''',
    '''Uyghur''': '''uig_Arab''',
    '''Ukrainian''': '''ukr_Cyrl''',
    '''Umbundu''': '''umb_Latn''',
    '''Urdu''': '''urd_Arab''',
    '''Northern Uzbek''': '''uzn_Latn''',
    '''Venetian''': '''vec_Latn''',
    '''Vietnamese''': '''vie_Latn''',
    '''Waray''': '''war_Latn''',
    '''Wolof''': '''wol_Latn''',
    '''Xhosa''': '''xho_Latn''',
    '''Eastern Yiddish''': '''ydd_Hebr''',
    '''Yoruba''': '''yor_Latn''',
    '''Yue Chinese''': '''yue_Hant''',
    '''Chinese Simplified''': '''zho_Hans''',
    '''Chinese Traditional''': '''zho_Hant''',
    '''Standard Malay''': '''zsm_Latn''',
    '''Zulu''': '''zul_Latn''',
}




class 			UpperCAmelCase					(      snake_case_       ):
      _lowercase:      Dict											=       '''facebook/nllb-200-distilled-600M'''
      _lowercase:      int											=       (
          '''This is a tool that translates text from a language to another. It takes three inputs: `text`, which should '''
          '''be the text to translate, `src_lang`, which should be the language of the text to translate and `tgt_lang`, '''
          '''which should be the language for the desired ouput language. Both `src_lang` and `tgt_lang` are written in '''
          '''plain English, such as \'Romanian\', or \'Albanian\'. It returns the text translated in `tgt_lang`.'''
      )
      _lowercase:      Any											=       '''translator'''
      _lowercase:      Optional[int]											=       AutoTokenizer
      _lowercase:      str											=       AutoModelForSeqaSeqLM
      _lowercase:      List[Any]											=       LANGUAGE_CODES

      _lowercase:      Tuple											=       ['''text''', '''text''', '''text''']
      _lowercase:      List[str]											=       ['''text''']




      def     lowercase__   (      self      :			str      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Optional[int]      ,      __snake_case      :			Dict  )    ->				Optional[Any]:
            if src_lang not in self.lang_to_code:
                  raise ValueError(f"{src_lang} is not a supported language."  )
            if tgt_lang not in self.lang_to_code:
                  raise ValueError(f"{tgt_lang} is not a supported language."  )
            _lowerCAmelCase												=							self.lang_to_code[src_lang]
            _lowerCAmelCase												=							self.lang_to_code[tgt_lang]
            return self.pre_processor._build_translation_inputs(
                __snake_case      ,      return_tensors="""pt"""      ,      src_lang=__snake_case      ,      tgt_lang=__snake_case  )




      def     lowercase__   (      self      :			Optional[int]      ,      __snake_case      :			Any  )    ->				List[str]:
            return self.model.generate(**__snake_case  )




      def     lowercase__   (      self      :			Dict      ,      __snake_case      :			List[Any]  )    ->				Any:
            return self.post_processor.decode(outputs[0].tolist()      ,      skip_special_tokens=__snake_case  )






