Chasing bug

tests/cli/test_neuron_cache_cli.py

@@ -25,7 +25,7 @@
 
 from huggingface_hub import HfApi, create_repo, delete_repo
 from huggingface_hub.utils import RepositoryNotFoundError
-from transformers import BertConfig, BertModel, BertTokenizer
+from transformers import BertConfig, BertModel, BertTokenizer, T5Tokenizer, T5Config, T5Model
 from transformers.testing_utils import is_staging_test
 
 from optimum.neuron.utils.cache_utils import (
@@ -115,15 +115,34 @@ def test_optimum_neuron_cache_add(self, hub_test):
 
             mandatory_tokens = ["[UNK]", "[SEP]", "[CLS]"]
 
-            with open(tmpdir / "vocab.txt", "w") as fp:
+            with open(tmpdir / "bert_vocab.txt", "w") as fp:
                 fp.write("\n".join([get_random_string(random.randint(10, 20))] + mandatory_tokens))
 
-            tokenizer = BertTokenizer(tmpdir / "vocab.txt")
+            tokenizer = BertTokenizer(tmpdir / "bert_vocab.txt")
             tokenizer.save_pretrained(bert_model_name)
 
             model = BertModel(config)
             model.save_pretrained(bert_model_name)
 
+            # Create dummy T5 model.
+            t5_model_name = tmpdir / "t5_model"
+            config = T5Config()
+
+            config.num_hidden_layers = 2
+            config.num_attention_heads = 2
+            config.vocab_size = 100
+
+            mandatory_tokens = ["[UNK]", "[SEP]", "[CLS]"]
+
+            with open(tmpdir / "t5_vocab.txt", "w") as fp:
+                fp.write("\n".join([get_random_string(random.randint(10, 20))] + mandatory_tokens))
+
+            tokenizer = T5Tokenizer(tmpdir / "t5_vocab.txt")
+            tokenizer.save_pretrained(t5_model_name)
+
+            model = T5Model(config)
+            model.save_pretrained(t5_model_name)
+
             env = dict(os.environ)
             env["OPTIMUM_NEURON_DISABLE_IS_PRIVATE_REPO_CHECK"] = "1"
 
@@ -158,7 +177,7 @@ def test_optimum_neuron_cache_add(self, hub_test):
 
             # seq2seq model.
             command = (
-                f"optimum-cli neuron cache add -m {bert_model_name} --task translation --train_batch_size 1 --precision bf16 "
+                f"optimum-cli neuron cache add -m {t5_model_name} --task translation --train_batch_size 1 --precision bf16 "
                 "--num_cores 2 --encoder_sequence_length 12 --decoder_sequence_length 12"
             ).split()
             p = subprocess.Popen(command, env=env)