Fix mistral gating (mosaicml#3199)

tests/common/models.py

@@ -651,32 +651,32 @@ def configure_tiny_t5_hf_model(use_logits: bool = True) -> HuggingFaceModel:
     return HuggingFaceModel(configure_tiny_t5_model(), configure_tiny_t5_tokenizer(), use_logits)
 
 
-def configure_tiny_mistral_model() -> 'PreTrainedModel':
+def configure_tiny_mpt_model() -> 'PreTrainedModel':
     try:
         from transformers import PreTrainedModel
-        assert isinstance(pytest.tiny_mistral_model, PreTrainedModel)
-        return copy.deepcopy(pytest.tiny_mistral_model)
+        assert isinstance(pytest.tiny_mpt_model, PreTrainedModel)
+        return copy.deepcopy(pytest.tiny_mpt_model)
     except AttributeError:
         pytest.skip('Composer installed without NLP support')
 
 
-def configure_tiny_mistral_tokenizer() -> Union['PreTrainedTokenizer', 'PreTrainedTokenizerFast']:
+def configure_tiny_mpt_tokenizer() -> Union['PreTrainedTokenizer', 'PreTrainedTokenizerFast']:
     try:
         from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
-        assert isinstance(pytest.tiny_mistral_tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast))
-        return copy.deepcopy(pytest.tiny_mistral_tokenizer)
+        assert isinstance(pytest.tiny_mpt_tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast))
+        return copy.deepcopy(pytest.tiny_mpt_tokenizer)
     except AttributeError:
         pytest.skip('Composer installed without NLP support')
 
 
-def configure_tiny_mistral_config() -> 'PretrainedConfig':
+def configure_tiny_mpt_config() -> 'PretrainedConfig':
     try:
         from transformers import PretrainedConfig
-        assert isinstance(pytest.tiny_mistral_config, PretrainedConfig)
-        return copy.deepcopy(pytest.tiny_mistral_config)
+        assert isinstance(pytest.tiny_mpt_config, PretrainedConfig)
+        return copy.deepcopy(pytest.tiny_mpt_config)
     except AttributeError:
         pytest.skip('Composer installed without NLP support')
 
 
-def configure_tiny_mistral_hf_model(use_logits: bool = True) -> HuggingFaceModel:
-    return HuggingFaceModel(configure_tiny_mistral_model(), configure_tiny_mistral_tokenizer(), use_logits)
+def configure_tiny_mpt_hf_model(use_logits: bool = True) -> HuggingFaceModel:
+    return HuggingFaceModel(configure_tiny_mpt_model(), configure_tiny_mpt_tokenizer(), use_logits)

tests/conftest.py

@@ -118,9 +118,9 @@ def pytest_configure():
             tiny_gpt2_model_helper,
             tiny_gpt2_tokenizer_helper,
             tiny_llama_tokenizer_helper,
-            tiny_mistral_config_helper,
-            tiny_mistral_model_helper,
-            tiny_mistral_tokenizer_helper,
+            tiny_mpt_config_helper,
+            tiny_mpt_model_helper,
+            tiny_mpt_tokenizer_helper,
             tiny_opt_config_helper,
             tiny_opt_model_helper,
             tiny_opt_tokenizer_helper,
@@ -141,9 +141,9 @@ def pytest_configure():
         pytest.tiny_t5_config = tiny_t5_config_helper()  # type: ignore
         pytest.tiny_t5_model = tiny_t5_model_helper(pytest.tiny_t5_config)  # type: ignore
         pytest.tiny_t5_tokenizer = tiny_t5_tokenizer_helper()  # type: ignore
-        pytest.tiny_mistral_config = tiny_mistral_config_helper()  # type: ignore
-        pytest.tiny_mistral_model = tiny_mistral_model_helper(pytest.tiny_mistral_config)  # type: ignore
-        pytest.tiny_mistral_tokenizer = tiny_mistral_tokenizer_helper()  # type: ignore
+        pytest.tiny_mpt_config = tiny_mpt_config_helper()  # type: ignore
+        pytest.tiny_mpt_model = tiny_mpt_model_helper(pytest.tiny_mpt_config)  # type: ignore
+        pytest.tiny_mpt_tokenizer = tiny_mpt_tokenizer_helper()  # type: ignore
 
 
 def pytest_sessionfinish(session: pytest.Session, exitstatus: int):

tests/fixtures/fixtures.py

@@ -333,45 +333,44 @@ def _session_tiny_t5_model(_session_tiny_t5_config):  # type: ignore
     return tiny_t5_model_helper(_session_tiny_t5_config)
 
 
-def tiny_mistral_config_helper():
+def tiny_mpt_config_helper():
     transformers = pytest.importorskip('transformers')
 
     tiny_overrides = {
-        'hidden_size': 128,
-        'intermediate_size': 256,
-        'num_attention_heads': 8,
-        'num_hidden_layers': 2,
-        'num_kv_heads': 4,
+        'd_model': 128,
+        'expansion_ratio': 1,
+        'n_heads': 8,
+        'n_layers': 2,
     }
-    return transformers.AutoConfig.from_pretrained('mistralai/Mistral-7B-v0.1', **tiny_overrides)
+    return transformers.AutoConfig.from_pretrained('mosaicml/mpt-7b', **tiny_overrides)
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_mistral_config():  # type: ignore
-    return tiny_mistral_config_helper()
+def _session_tiny_mpt_config():  # type: ignore
+    return tiny_mpt_config_helper()
 
 
-def tiny_mistral_tokenizer_helper():
+def tiny_mpt_tokenizer_helper():
     transformers = pytest.importorskip('transformers')
 
-    hf_tokenizer = transformers.AutoTokenizer.from_pretrained('mistralai/Mistral-7B-v0.1', model_max_length=512)
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained('mosaicml/mpt-7b', model_max_length=512)
     return hf_tokenizer
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_mistral_tokenizer():  # type: ignore
-    return tiny_mistral_tokenizer_helper()
+def _session_tiny_mpt_tokenizer():  # type: ignore
+    return tiny_mpt_tokenizer_helper()
 
 
-def tiny_mistral_model_helper(config):
+def tiny_mpt_model_helper(config):
     transformers = pytest.importorskip('transformers')
 
     return transformers.AutoModelForCausalLM.from_config(config)
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_mistral_model(_session_tiny_mistral_config):  # type: ignore
-    return tiny_mistral_model_helper(_session_tiny_mistral_config)
+def _session_tiny_mpt_model(_session_tiny_mpt_config):  # type: ignore
+    return tiny_mpt_model_helper(_session_tiny_mpt_config)
 
 
 @pytest.fixture
@@ -455,15 +454,15 @@ def tiny_t5_model(_session_tiny_t5_model):
 
 
 @pytest.fixture
-def tiny_mistral_config(_session_tiny_mistral_config):
-    return copy.deepcopy(_session_tiny_mistral_config)
+def tiny_mpt_config(_session_tiny_mpt_config):
+    return copy.deepcopy(_session_tiny_mpt_config)
 
 
 @pytest.fixture
-def tiny_mistral_tokenizer(_session_tiny_mistral_tokenizer):
-    return copy.deepcopy(_session_tiny_mistral_tokenizer)
+def tiny_mpt_tokenizer(_session_tiny_mpt_tokenizer):
+    return copy.deepcopy(_session_tiny_mpt_tokenizer)
 
 
 @pytest.fixture
-def tiny_mistral_model(_session_tiny_mistral_model):
-    return copy.deepcopy(_session_tiny_mistral_model)
+def tiny_mpt_model(_session_tiny_mpt_model):
+    return copy.deepcopy(_session_tiny_mpt_model)

tests/loggers/test_mlflow_logger.py

@@ -362,8 +362,7 @@ def test_mlflow_log_model(tmp_path, tiny_gpt2_model, tiny_gpt2_tokenizer):
             'tokenizer': tiny_gpt2_tokenizer,
         },
         artifact_path='my_model',
-        metadata={'task': 'llm/v1/completions'},
-        task='text-generation',
+        task='llm/v1/completions',
     )
     test_mlflow_logger.post_close()
 
@@ -405,8 +404,7 @@ def test_mlflow_save_model(tmp_path, tiny_gpt2_model, tiny_gpt2_tokenizer):
             'tokenizer': tiny_gpt2_tokenizer,
         },
         path=local_mlflow_save_path,
-        metadata={'task': 'llm/v1/completions'},
-        task='text-generation',
+        task='llm/v1/completions',
     )
     test_mlflow_logger.post_close()
 
@@ -418,16 +416,16 @@ def test_mlflow_save_model(tmp_path, tiny_gpt2_model, tiny_gpt2_tokenizer):
 
 @pytest.mark.filterwarnings('ignore:.*Setuptools is replacing distutils.*:UserWarning')
 @pytest.mark.filterwarnings("ignore:.*The 'transformers' MLflow Models integration.*:FutureWarning")
-def test_mlflow_save_peft_model(tmp_path, tiny_mistral_model, tiny_mistral_tokenizer):
+def test_mlflow_save_peft_model(tmp_path, tiny_mpt_model, tiny_mpt_tokenizer):
     mlflow = pytest.importorskip('mlflow')
     peft = pytest.importorskip('peft')
 
     # Reload just so the model has the update base model name
-    tiny_mistral_model.save_pretrained(tmp_path / Path('tiny_mistral_save_pt'))
-    tiny_mistral_model = tiny_mistral_model.from_pretrained(tmp_path / Path('tiny_mistral_save_pt'))
+    tiny_mpt_model.save_pretrained(tmp_path / Path('tiny_mpt_save_pt'))
+    tiny_mpt_model = tiny_mpt_model.from_pretrained(tmp_path / Path('tiny_mpt_save_pt'))
 
     peft_config = {'peft_type': 'LORA'}
-    peft_model = peft.get_peft_model(tiny_mistral_model, peft.get_peft_config(peft_config))
+    peft_model = peft.get_peft_model(tiny_mpt_model, peft.get_peft_config(peft_config))
 
     mlflow_uri = tmp_path / Path('my-test-mlflow-uri')
     mlflow_exp_name = 'test-log-model-exp-name'
@@ -441,7 +439,7 @@ def test_mlflow_save_peft_model(tmp_path, tiny_mistral_model, tiny_mistral_token
     mock_logger = MagicMock()
 
     peft_model.save_pretrained(tmp_path / Path('peft_model_save_pt'))
-    tiny_mistral_tokenizer.save_pretrained(tmp_path / Path('peft_model_save_pt'))
+    tiny_mpt_tokenizer.save_pretrained(tmp_path / Path('peft_model_save_pt'))
 
     local_mlflow_save_path = str(tmp_path / Path('my_model_local'))
     test_mlflow_logger.init(state=mock_state, logger=mock_logger)
@@ -454,8 +452,8 @@ def test_mlflow_save_peft_model(tmp_path, tiny_mistral_model, tiny_mistral_token
 
     loaded_model = mlflow.pyfunc.load_model(local_mlflow_save_path).unwrap_python_model()
 
-    check_hf_model_equivalence(loaded_model.model, tiny_mistral_model)
-    check_hf_tokenizer_equivalence(loaded_model.tokenizer, tiny_mistral_tokenizer)
+    check_hf_model_equivalence(loaded_model.model, tiny_mpt_model)
+    check_hf_tokenizer_equivalence(loaded_model.tokenizer, tiny_mpt_tokenizer)
 
 
 @pytest.mark.filterwarnings('ignore:.*Setuptools is replacing distutils.*:UserWarning')
@@ -688,6 +686,7 @@ class ImageLogger(Callback):
         def before_forward(self, state: State, logger: Logger):
             inputs = state.batch_get_item(key=0)
             images = inputs.data.cpu().numpy()
+            images = np.clip(images, 0, 1)
             logger.log_images(images, step=state.timestamp.batch.value)
             with pytest.warns(UserWarning):
                 logger.log_images(

tests/models/test_hf_model.py

@@ -30,8 +30,8 @@
     configure_tiny_bert_tokenizer,
     configure_tiny_gpt2_model,
     configure_tiny_gpt2_tokenizer,
-    configure_tiny_mistral_model,
-    configure_tiny_mistral_tokenizer,
+    configure_tiny_mpt_model,
+    configure_tiny_mpt_tokenizer,
     configure_tiny_t5_model,
     configure_tiny_t5_tokenizer,
 )
@@ -59,7 +59,7 @@ def gpt2_peft_config():
     return _gpt2_peft_config()
 
 
-def _mistral_peft_config():
+def _mpt_peft_config():
     pytest.importorskip('peft')
     from peft import get_peft_config
 
@@ -72,8 +72,8 @@ def _mistral_peft_config():
 
 
 @pytest.fixture
-def mistral_peft_config():
-    return _mistral_peft_config()
+def mpt_peft_config():
+    return _mpt_peft_config()
 
 
 def test_hf_tokenizer_save(tmp_path: Path, tiny_bert_model, tiny_bert_tokenizer):
@@ -1422,7 +1422,7 @@ def test_peft_trains_and_loads(tiny_gpt2_model, tiny_gpt2_tokenizer, gpt2_peft_c
     'model,tokenizer,peft_config',
     [
         (configure_tiny_gpt2_model, configure_tiny_gpt2_tokenizer, _gpt2_peft_config()),
-        (configure_tiny_mistral_model, configure_tiny_mistral_tokenizer, _mistral_peft_config()),
+        (configure_tiny_mpt_model, configure_tiny_mpt_tokenizer, _mpt_peft_config()),
     ],
 )
 def test_peft_generate(model, tokenizer, peft_config):