
import argparse
import logging
import os
import sys
import tempfile
from pathlib import Path

import lightning_base
import pytest
import pytorch_lightning as pl
import torch
from convert_pl_checkpoint_to_hf import convert_pl_to_hf
from distillation import distill_main
from finetune import SummarizationModule, main
from huggingface_hub import list_models
from parameterized import parameterized
from run_eval import generate_summaries_or_translations
from torch import nn

from transformers import AutoConfig, AutoModelForSeqaSeqLM
from transformers.testing_utils import CaptureStderr, CaptureStdout, TestCasePlus, require_torch_gpu, slow
from utils import label_smoothed_nll_loss, lmap, load_json


logging.basicConfig(level=logging.DEBUG)

_UpperCAmelCase    			=					logging.getLogger()
_UpperCAmelCase    			=					torch.cuda.is_available()
_UpperCAmelCase    			=					{
    """max_tokens_per_batch""": None,
    """supervise_forward""": True,
    """normalize_hidden""": True,
    """label_smoothing""": 0.2,
    """eval_max_gen_length""": None,
    """eval_beams""": 1,
    """val_metric""": """loss""",
    """save_top_k""": 1,
    """adafactor""": True,
    """early_stopping_patience""": 2,
    """logger_name""": """default""",
    """length_penalty""": 0.5,
    """cache_dir""": """""",
    """task""": """summarization""",
    """num_workers""": 2,
    """alpha_hid""": 0,
    """freeze_embeds""": True,
    """enc_only""": False,
    """tgt_suffix""": """""",
    """resume_from_checkpoint""": None,
    """sortish_sampler""": True,
    """student_decoder_layers""": 1,
    """val_check_interval""": 1.0,
    """output_dir""": """""",
    """fp16""": False,  # TODO(SS): set this to CUDA_AVAILABLE if ci installs apex or start using native amp
    """no_teacher""": False,
    """fp16_opt_level""": """O1""",
    """gpus""": 1 if CUDA_AVAILABLE else 0,
    """n_tpu_cores""": 0,
    """max_grad_norm""": 1.0,
    """do_train""": True,
    """do_predict""": True,
    """accumulate_grad_batches""": 1,
    """server_ip""": """""",
    """server_port""": """""",
    """seed""": 42,
    """model_name_or_path""": """sshleifer/bart-tiny-random""",
    """config_name""": """""",
    """tokenizer_name""": """facebook/bart-large""",
    """do_lower_case""": False,
    """learning_rate""": 0.3,
    """lr_scheduler""": """linear""",
    """weight_decay""": 0.0,
    """adam_epsilon""": 1e-08,
    """warmup_steps""": 0,
    """max_epochs""": 1,
    """train_batch_size""": 2,
    """eval_batch_size""": 2,
    """max_source_length""": 12,
    """max_target_length""": 12,
    """val_max_target_length""": 12,
    """test_max_target_length""": 12,
    """fast_dev_run""": False,
    """no_cache""": False,
    """n_train""": -1,
    """n_val""": -1,
    """n_test""": -1,
    """student_encoder_layers""": 1,
    """freeze_encoder""": False,
    """auto_scale_batch_size""": False,
    """overwrite_output_dir""": False,
    """student""": None,
}


def 		_lowerCamelCase       (   _a			,   _a  ):
  """simple docstring"""




  _lowerCamelCase      =						"\n".join(_a  )
  Path(_a  ).open('''w'''  ).writelines(_a  )


_UpperCAmelCase    			=					[""" Sam ate lunch today.""", """Sams lunch ingredients."""]
_UpperCAmelCase    			=					["""A very interesting story about what I ate for lunch.""", """Avocado, celery, turkey, coffee"""]
_UpperCAmelCase    			=					"""patrickvonplaten/t5-tiny-random"""
_UpperCAmelCase    			=					"""sshleifer/t5-tinier-random"""
_UpperCAmelCase    			=					"""sshleifer/bart-tiny-random"""
_UpperCAmelCase    			=					"""sshleifer/tiny-mbart"""
_UpperCAmelCase    			=					"""sshleifer/tiny-marian-en-de"""
_UpperCAmelCase    			=					"""stas/tiny-wmt19-en-de"""


_UpperCAmelCase    			=					logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks


def 		_lowerCamelCase       (   _a  ):
  """simple docstring"""




  for split in ["train", "val", "test"]:
    _dump_articles(os.path.join(_a			,   F'''{split}.source'''  )			,   _a  )
    _dump_articles(os.path.join(_a			,   F'''{split}.target'''  )			,   _a  )
  return tmp_dir

class 		__magic_name__				( lowerCAmelCase_					):


  """simple docstring"""





  @classmethod
  def     _UpperCAmelCase  (					cls		):
    logging.disable(logging.CRITICAL		)  # remove noisy download output from tracebacks
    return cls


  @slow
  @require_torch_gpu
  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						list_models()
    _lowerCamelCase      =						"sshleifer"
    _lowerCamelCase      =						[x.modelId for x in model_list if x.modelId.startswith(__UpperCamelCase		)]
    _lowerCamelCase      =						["sshleifer/blenderbot-3B", "sshleifer/blenderbot-90M"]
    _lowerCamelCase      =						[]
    for m in model_ids:
      if m in allowed_to_be_broken:
        continue
      try:
        AutoConfig.from_pretrained(__UpperCamelCase		)
      except Exception:
        failures.append(__UpperCamelCase		)
    assert not failures, f'''The following models could not be loaded through AutoConfig: {failures}'''


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						{"student_encoder_layers": 2, "student_decoder_layers": 1, "no_teacher": True}
    self._test_distiller_cli(__UpperCamelCase		)


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						{
        "student_encoder_layers": 2,
        "student_decoder_layers": 1,
        "max_epochs": 4,
        "val_check_interval": 0.25,
        "alpha_hid": 2.0,
        "model_name_or_path": "IGNORE_THIS_IT_DOESNT_GET_USED",
    }
    _lowerCamelCase      =						self._test_distiller_cli(__UpperCamelCase	,		check_contents=__UpperCamelCase		)

    _lowerCamelCase      =						list(Path(model.output_dir		).glob('''*.ckpt'''		)		)
    self.assertEqual(1	,		len(__UpperCamelCase		)		)
    _lowerCamelCase      =						list(Path(model.output_dir		).glob('''**/*.bin'''		)		)
    self.assertEqual(len(__UpperCamelCase		)	,		2		)
    _lowerCamelCase      =						lmap(str.strip	,		Path(model.hparams.data_dir		).joinpath('''test.source'''		).open().readlines()		)
    _lowerCamelCase      =						tempfile.mktemp()  # XXX: not being cleaned up
    generate_summaries_or_translations(__UpperCamelCase	,		__UpperCamelCase	,		str(model.output_dir / '''best_tfmr'''		)		)
    self.assertTrue(Path(__UpperCamelCase		).exists()		)

    _lowerCamelCase      =						self.get_auto_remove_tmp_dir()
    convert_pl_to_hf(ckpts[0]	,		transformer_ckpts[0].parent	,		__UpperCamelCase		)
    assert os.path.exists(os.path.join(__UpperCamelCase	,		'''pytorch_model.bin'''		)		)


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						AutoModelForSeqaSeqLM.from_pretrained(__UpperCamelCase		)
    _lowerCamelCase      =						model.dummy_inputs["input_ids"], model.dummy_inputs["attention_mask"]
    _lowerCamelCase      =						torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]]	,		dtype=torch.long	,		device=model.device		)
    _lowerCamelCase      =						target_ids[:, :-1].contiguous()  # Why this line?
    _lowerCamelCase      =						target_ids[:, 1:].clone()  # why clone?
    _lowerCamelCase      =						model(
        __UpperCamelCase	,		attention_mask=__UpperCamelCase	,		decoder_input_ids=__UpperCamelCase	,		labels=__UpperCamelCase	,		use_cache=__UpperCamelCase		).loss

    _lowerCamelCase      =						model(__UpperCamelCase	,		attention_mask=__UpperCamelCase	,		decoder_input_ids=__UpperCamelCase	,		use_cache=__UpperCamelCase		).logits

    _lowerCamelCase      =						nn.functional.log_softmax(__UpperCamelCase	,		dim=-1		)
    _lowerCamelCase      =						label_smoothed_nll_loss(
        __UpperCamelCase	,		__UpperCamelCase	,		0.1	,		ignore_index=model.config.pad_token_id		)
    with self.assertRaises(__UpperCamelCase		):
      # TODO: understand why this breaks
      self.assertEqual(__UpperCamelCase	,		__UpperCamelCase		)


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						{
        "student_encoder_layers": 2,
        "student_decoder_layers": 1,
        "num_train_epochs": 4,
        "val_check_interval": 0.25,
        "alpha_hid": 2.0,
        "task": "translation",
        "model_name_or_path": "IGNORE_THIS_IT_DOESNT_GET_USED",
        "tokenizer_name": MBART_TINY,
        "teacher": MBART_TINY,
        "src_lang": "en_XX",
        "tgt_lang": "ro_RO",
    }
    _lowerCamelCase      =						self._test_distiller_cli(__UpperCamelCase	,		check_contents=__UpperCamelCase		)
    assert model.model.config.model_type == "mbart"

    _lowerCamelCase      =						list(Path(model.output_dir		).glob('''*.ckpt'''		)		)
    self.assertEqual(1	,		len(__UpperCamelCase		)		)
    _lowerCamelCase      =						list(Path(model.output_dir		).glob('''**/*.bin'''		)		)
    _lowerCamelCase      =						list(Path(model.output_dir		).glob('''best_tfmr/*'''		)		)
    assert len(__UpperCamelCase		) > 2
    self.assertEqual(len(__UpperCamelCase		)	,		2		)


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						{
        "student_encoder_layers": 1,
        "student_decoder_layers": 1,
        "alpha_hid": 2.0,
        "teacher": T5_TINY,
        "model_name_or_path": T5_TINY,
        "tokenizer_name": T5_TINY,
    }
    self._test_distiller_cli(__UpperCamelCase		)


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						{
        "teacher": T5_TINY,
        "student": T5_TINIER,
        "model_name_or_path": T5_TINIER,
        "tokenizer_name": T5_TINIER,
    }
    self._test_distiller_cli(__UpperCamelCase		)


  def     _UpperCAmelCase  (					self	,		a__	,		a__=True		):
    _lowerCamelCase      =						{
        "label_smoothing": 0.0,
        "early_stopping_patience": -1,
        "train_batch_size": 1,
        "eval_batch_size": 2,
        "max_epochs": 2,
        "alpha_mlm": 0.2,
        "alpha_ce": 0.8,
        "do_predict": True,
        "model_name_or_path": "sshleifer/tinier_bart",
        "teacher": CHEAP_ARGS["model_name_or_path"],
        "val_check_interval": 0.5,
    }
    default_updates.update(__UpperCamelCase		)
    _lowerCamelCase      =						CHEAP_ARGS.copy()
    _lowerCamelCase      =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()		)
    _lowerCamelCase      =						self.get_auto_remove_tmp_dir()

    args_d.update(data_dir=__UpperCamelCase	,		output_dir=__UpperCamelCase	,		**__UpperCamelCase		)
    _lowerCamelCase      =						distill_main(argparse.Namespace(**__UpperCamelCase		)		)
    if not check_contents:
      return model
    _lowerCamelCase      =						os.listdir(__UpperCamelCase		)
    _lowerCamelCase      =						{os.path.basename(__UpperCamelCase		) for p in contents}
    _lowerCamelCase      =						[p for p in contents if p.endswith('''ckpt'''		)]
    assert len(__UpperCamelCase		) > 0

    self.assertIn('''test_generations.txt'''	,		__UpperCamelCase		)
    self.assertIn('''test_results.txt'''	,		__UpperCamelCase		)

    _lowerCamelCase      =						load_json(model.metrics_save_path		)
    _lowerCamelCase      =						metrics["val"][-1]
    self.assertGreaterEqual(last_step_stats['''val_avg_gen_time''']	,		0.01		)
    self.assertGreaterEqual(1.0	,		last_step_stats['''val_avg_gen_time''']		)
    self.assertIsInstance(last_step_stats[f'''val_avg_{model.val_metric}''']	,		__UpperCamelCase		)
    _lowerCamelCase      =						int(args_d['''max_epochs'''] * (1 / args_d['''val_check_interval''']) + 1		)
    self.assertEqual(len(metrics['''val''']		)	,		__UpperCamelCase		)
    self.assertEqual(len(metrics['''test''']		)	,		1		)
    return model







class 		__magic_name__				( lowerCAmelCase_					):


  """simple docstring"""





  @parameterized.expand(
      [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY]	,		)
  def     _UpperCAmelCase  (					self	,		a__		):
    _lowerCamelCase      =						CHEAP_ARGS.copy()
    _lowerCamelCase      =						"translation" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else "summarization"
    _lowerCamelCase      =						0.1 if task == "translation" else 0

    _lowerCamelCase      =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()		)
    _lowerCamelCase      =						self.get_auto_remove_tmp_dir()
    args_d.update(
        data_dir=__UpperCamelCase	,		model_name_or_path=__UpperCamelCase	,		tokenizer_name=__UpperCamelCase	,		train_batch_size=2	,		eval_batch_size=2	,		output_dir=__UpperCamelCase	,		do_predict=__UpperCamelCase	,		task=__UpperCamelCase	,		src_lang='''en_XX'''	,		tgt_lang='''ro_RO'''	,		freeze_encoder=__UpperCamelCase	,		freeze_embeds=__UpperCamelCase	,		)
    assert "n_train" in args_d
    _lowerCamelCase      =						argparse.Namespace(**__UpperCamelCase		)
    _lowerCamelCase      =						main(__UpperCamelCase		)

    _lowerCamelCase      =						module.model.get_input_embeddings()
    assert not input_embeds.weight.requires_grad
    if model == T5_TINY:
      _lowerCamelCase      =						module.model.lm_head
      assert not lm_head.weight.requires_grad
      assert (lm_head.weight == input_embeds.weight).all().item()
    elif model == FSMT_TINY:
      _lowerCamelCase      =						module.model.model
      _lowerCamelCase      =						fsmt.decoder.embed_positions
      assert not embed_pos.weight.requires_grad
      assert not fsmt.decoder.embed_tokens.weight.requires_grad
      # check that embeds are not the same
      assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
    else:
      _lowerCamelCase      =						module.model.model
      _lowerCamelCase      =						bart.decoder.embed_positions
      assert not embed_pos.weight.requires_grad
      assert not bart.shared.weight.requires_grad
      # check that embeds are the same
      assert bart.decoder.embed_tokens == bart.encoder.embed_tokens
      assert bart.decoder.embed_tokens == bart.shared

    _lowerCamelCase      =						load_json(module.output_dir / '''text_batch.json'''		)
    assert isinstance(__UpperCamelCase	,		__UpperCamelCase		)
    assert len(__UpperCamelCase		) >= 4


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						CHEAP_ARGS.copy()

    _lowerCamelCase      =						"summarization"
    _lowerCamelCase      =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()		)

    args_d.update(
        data_dir=__UpperCamelCase	,		tokenizer_name=__UpperCamelCase	,		train_batch_size=2	,		eval_batch_size=2	,		do_predict=__UpperCamelCase	,		task=__UpperCamelCase	,		src_lang='''en_XX'''	,		tgt_lang='''ro_RO'''	,		freeze_encoder=__UpperCamelCase	,		freeze_embeds=__UpperCamelCase	,		)

    # test models whose config includes the extra_model_args
    _lowerCamelCase      =						BART_TINY
    _lowerCamelCase      =						self.get_auto_remove_tmp_dir()
    _lowerCamelCase      =						args_d.copy()
    args_da.update(
        model_name_or_path=__UpperCamelCase	,		output_dir=__UpperCamelCase	,		)
    _lowerCamelCase      =						("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
    for p in extra_model_params:
      _lowerCamelCase      =						0.5
    _lowerCamelCase      =						argparse.Namespace(**__UpperCamelCase		)
    _lowerCamelCase      =						main(__UpperCamelCase		)
    for p in extra_model_params:
      assert getattr(model.config	,		__UpperCamelCase		) == 0.5, f'''failed to override the model config for param {p}'''

    # test models whose config doesn't include the extra_model_args
    _lowerCamelCase      =						T5_TINY
    _lowerCamelCase      =						self.get_auto_remove_tmp_dir()
    _lowerCamelCase      =						args_d.copy()
    args_da.update(
        model_name_or_path=__UpperCamelCase	,		output_dir=__UpperCamelCase	,		)
    _lowerCamelCase      =						"encoder_layerdrop"
    _lowerCamelCase      =						0.5
    _lowerCamelCase      =						argparse.Namespace(**__UpperCamelCase		)
    with pytest.raises(__UpperCamelCase		) as excinfo:
      _lowerCamelCase      =						main(__UpperCamelCase		)
    assert str(excinfo.value		) == f'''model config doesn\'t have a `{unsupported_param}` attribute'''


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						CHEAP_ARGS.copy()

    _lowerCamelCase      =						"summarization"
    _lowerCamelCase      =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()		)

    _lowerCamelCase      =						BART_TINY
    _lowerCamelCase      =						self.get_auto_remove_tmp_dir()

    args_d.update(
        data_dir=__UpperCamelCase	,		model_name_or_path=__UpperCamelCase	,		output_dir=__UpperCamelCase	,		tokenizer_name=__UpperCamelCase	,		train_batch_size=2	,		eval_batch_size=2	,		do_predict=__UpperCamelCase	,		task=__UpperCamelCase	,		src_lang='''en_XX'''	,		tgt_lang='''ro_RO'''	,		freeze_encoder=__UpperCamelCase	,		freeze_embeds=__UpperCamelCase	,		)

    # emulate finetune.py
    _lowerCamelCase      =						argparse.ArgumentParser()
    _lowerCamelCase      =						pl.Trainer.add_argparse_args(__UpperCamelCase		)
    _lowerCamelCase      =						SummarizationModule.add_model_specific_args(__UpperCamelCase	,		os.getcwd()		)
    _lowerCamelCase      =						{"--help": True}

    # --help test
    with pytest.raises(__UpperCamelCase		) as excinfo:
      with CaptureStdout() as cs:
        _lowerCamelCase      =						parser.parse_args(__UpperCamelCase		)
      assert False, "--help is expected to sys.exit"
    assert excinfo.type == SystemExit
    _lowerCamelCase      =						lightning_base.arg_to_scheduler_metavar
    assert expected in cs.out, "--help is expected to list the supported schedulers"

    # --lr_scheduler=non_existing_scheduler test
    _lowerCamelCase      =						"non_existing_scheduler"
    _lowerCamelCase      =						{f'''--lr_scheduler={unsupported_param}'''}
    with pytest.raises(__UpperCamelCase		) as excinfo:
      with CaptureStderr() as cs:
        _lowerCamelCase      =						parser.parse_args(__UpperCamelCase		)
      assert False, "invalid argument is expected to sys.exit"
    assert excinfo.type == SystemExit
    _lowerCamelCase      =						f'''invalid choice: \'{unsupported_param}\''''
    assert expected in cs.err, f'''should have bailed on invalid choice of scheduler {unsupported_param}'''

    # --lr_scheduler=existing_scheduler test
    _lowerCamelCase      =						"cosine"
    _lowerCamelCase      =						args_d.copy()
    _lowerCamelCase      =						supported_param
    _lowerCamelCase      =						argparse.Namespace(**__UpperCamelCase		)
    _lowerCamelCase      =						main(__UpperCamelCase		)
    assert (
        getattr(model.hparams	,		'''lr_scheduler'''		) == supported_param
    ), f'''lr_scheduler={supported_param} shouldn\'t fail'''






"""simple docstring"""




import argparse
import os


A__   :						Any 							=     {
    """pytorch""": {
        """1.13""": {
            """torch""": """1.13.1""",
            """torchvision""": """0.14.1""",
            """torchaudio""": """0.13.1""",
            """python""": 3.9,
            """cuda""": """cu116""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1"""
                """ --extra-index-url https://download.pytorch.org/whl/cu116"""
            ),
            """base_image""": """nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04""",
        },
        """1.12""": {
            """torch""": """1.12.1""",
            """torchvision""": """0.13.1""",
            """torchaudio""": """0.12.1""",
            """python""": 3.9,
            """cuda""": """cu113""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1"""
                """ --extra-index-url https://download.pytorch.org/whl/cu113"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """1.11""": {
            """torch""": """1.11.0""",
            """torchvision""": """0.12.0""",
            """torchaudio""": """0.11.0""",
            """python""": 3.9,
            """cuda""": """cu113""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0"""
                """ --extra-index-url https://download.pytorch.org/whl/cu113"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """1.10""": {
            """torch""": """1.10.2""",
            """torchvision""": """0.11.3""",
            """torchaudio""": """0.10.2""",
            """python""": 3.9,
            """cuda""": """cu113""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2"""
                """ --extra-index-url https://download.pytorch.org/whl/cu113"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        # torchaudio < 0.10 has no CUDA-enabled binary distributions
        """1.9""": {
            """torch""": """1.9.1""",
            """torchvision""": """0.10.1""",
            """torchaudio""": """0.9.1""",
            """python""": 3.9,
            """cuda""": """cu111""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1"""
                """ --extra-index-url https://download.pytorch.org/whl/cu111"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
    },
    """tensorflow""": {
        """2.11""": {
            """tensorflow""": """2.11.1""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.11.1""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.10""": {
            """tensorflow""": """2.10.1""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.10.1""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.9""": {
            """tensorflow""": """2.9.3""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.9.3""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.8""": {
            """tensorflow""": """2.8.2""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.8.2""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.7""": {
            """tensorflow""": """2.7.3""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.7.3""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.6""": {
            """tensorflow""": """2.6.5""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.6.5""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.5""": {
            """tensorflow""": """2.5.3""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.5.3""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
    },
}


if __name__ == "__main__":
    A__   :						str 							=     argparse.ArgumentParser("""Choose the framework and version to install""")
    parser.add_argument(
        """--framework""", help="""The framework to install. Should be `torch` or `tensorflow`""", type=str, required=True
    )
    parser.add_argument("""--version""", help="""The version of the framework to install.""", type=str, required=True)
    A__   :						Any 							=     parser.parse_args()

    A__   :						List[Any] 							=     past_versions_testing[args.framework][args.version]

    os.system(f"""echo \"export INSTALL_CMD=\'{info["install"]}\'\" >> ~/.profile""")
    print(f"""echo \"export INSTALL_CMD=\'{info["install"]}\'\" >> ~/.profile""")

    A__   :						Any 							=     """"""
    if args.framework == "pytorch":
        A__   :						str 							=     info["""cuda"""]
    os.system(f"""echo \"export CUDA='{cuda}'\" >> ~/.profile""")
    print(f"""echo \"export CUDA='{cuda}'\" >> ~/.profile""")



