Repository: huggingface/transformers
File: tests/utils/test_modeling_utils.py (Python)

Code before:
        model = AutoModel.from_config(config, dtype=torch.float16)
        self.assertEqual(model.dtype, torch.float16)

        # torch.set_default_dtype() supports only float dtypes, so will fail with non-float type
        with self.assertRaises(ValueError):
            model = AutoModel.from_config(config, dtype=torch.int64)

    def test_model_from_config_dtype_str(self):
        # test that from_pretrained works with dtype being strings like "float32" for PyTorch backend
        model = AutoModel.from_pretrained(TINY_T5, dtype="float32")
        self.assertEqual(model.dtype, torch.float32)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        model = AutoModel.from_pretrained(TINY_T5, dtype="float16")
        self.assertEqual(model.dtype, torch.float16)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        # torch.set_default_dtype() supports only float dtypes, so will fail with non-float type
        with self.assertRaises(ValueError):
            model = AutoModel.from_pretrained(TINY_T5, dtype="int64")

    def test_model_from_config_dtype_composite(self):
        """
        Test that from_pretrained works with dtype being as a dict per each sub-config in composite config
        Tiny-Llava has saved auto dtype as `torch.float32` for all modules.
        Note, this is a deprecated feature and we fallbacl to main dtype in all cases below. This tests checks
        if the dtype fallback works correctly.
        """
        # Load without dtype specified
        model = LlavaForConditionalGeneration.from_pretrained(TINY_LLAVA)
        self.assertEqual(model.model.language_model.dtype, torch.float32)
        self.assertEqual(model.model.vision_tower.dtype, torch.float32)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        # should be able to set dtype as a simple string and the model loads it correctly
        model = LlavaForConditionalGeneration.from_pretrained(TINY_LLAVA, dtype="float32")
        self.assertEqual(model.model.language_model.dtype, torch.float32)
        self.assertEqual(model.model.vision_tower.dtype, torch.float32)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        model = LlavaForConditionalGeneration.from_pretrained(TINY_LLAVA, dtype=torch.float16)
        self.assertEqual(model.model.language_model.dtype, torch.float16)
        self.assertEqual(model.model.vision_tower.dtype, torch.float16)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        # should be able to accept dtype as a dict for each sub-config
        model = LlavaForConditionalGeneration.from_pretrained(
            TINY_LLAVA, dtype={"text_config": "float32", "vision_config": "float16", "": "bfloat16"}
        )
        self.assertEqual(model.model.language_model.dtype, torch.bfloat16)
        self.assertEqual(model.model.vision_tower.dtype, torch.bfloat16)

Review comment (suggestion):
Nit:
```suggestion
        Note, this is a deprecated feature and we fallback to main dtype in all cases below. This test checks
```

Code after:
        # Needs a device_map for to enter the low_cpu_mem branch. We also load AutoModelForSequenceClassification
        # deliberately to enter the missing keys branch.
        model = AutoModelForSequenceClassification.from_pretrained(
            TINY_MISTRAL, device_map="auto", quantization_config=None
        )
        self.assertIsNotNone(model)

    def test_model_from_config_dtype(self):
        # test that the model can be instantiated with dtype of user's choice - as long as it's a
        # float dtype. To make it happen config.dtype needs to be set before instantiating the
        # model from the config object.

        config = T5Config.from_pretrained(TINY_T5)
        model = AutoModel.from_config(config)
        # XXX: isn't supported
        # model = T5ForConditionalGeneration.from_config(config)
        self.assertEqual(model.dtype, torch.float32)

        model = AutoModel.from_config(config, dtype=torch.float16)
        self.assertEqual(model.dtype, torch.float16)

        # torch.set_default_dtype() supports only float dtypes, so will fail with non-float type
        with self.assertRaises(ValueError):
            model = AutoModel.from_config(config, dtype=torch.int64)

    def test_model_from_config_dtype_str(self):
        # test that from_pretrained works with dtype being strings like "float32" for PyTorch backend
        model = AutoModel.from_pretrained(TINY_T5, dtype="float32")
        self.assertEqual(model.dtype, torch.float32)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        model = AutoModel.from_pretrained(TINY_T5, dtype="float16")
        self.assertEqual(model.dtype, torch.float16)
        self.assertIsInstance(model.config.dtype, torch.dtype)

        # torch.set_default_dtype() supports only float dtypes, so will fail with non-float type
        with self.assertRaises(ValueError):
            model = AutoModel.from_pretrained(TINY_T5, dtype="int64")

    def test_model_from_config_dtype_composite(self):
        """
        Test that from_pretrained works with dtype being as a dict per each sub-config in composite config
        Tiny-Llava has saved auto dtype as `torch.float32` for all modules.
        Note, this is a deprecated feature and we fallback to main dtype in all cases below. This test checks
        if the dtype fallback works correctly.
        """
        # Load without dtype specified
        model = LlavaForConditionalGeneration.from_pretrained(TINY_LLAVA)
        self.assertEqual(model.model.language_model.dtype, torch.float32)
        self.assertEqual(model.model.vision_tower.dtype, torch.float32)
        self.assertIsInstance(model.config.dtype, torch.dtype)