Repository: huggingface/transformers
File: tests/models/jais2/test_modeling_jais2.py (Python)

Code before:

    @require_deterministic_for_xpu
    def test_model_logits(self):
        model_id = "inceptionai/Jais-2-8B-Chat"
        dummy_input = torch.LongTensor([[0, 0, 0, 0, 0, 0, 1, 2, 3], [1, 1, 2, 3, 4, 5, 6, 7, 8]]).to(torch_device)
        attention_mask = dummy_input.ne(0).to(torch.long)

        model = Jais2ForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")

        with torch.no_grad():
            logits = model(dummy_input, attention_mask=attention_mask).logits
        logits = logits.float()

        # fmt: off
        expected_logits_batch0 = Expectations(
            {
                ("cuda", None): [-0.9751, -1.0918, -0.9600, -0.9526, -0.9600, -0.9551, -0.9624, -0.9644, -0.9644, -0.9600, -0.9561, -0.9658, -0.9585, -0.9688, -0.9663],
                ("xpu", 3): [-0.9692, -1.0859, -0.9541, -0.9468, -0.9546, -0.9492, -0.9570, -0.9585, -0.9585, -0.9541, -0.9507, -0.9604, -0.9526, -0.9634, -0.9609],
            }
        )
        expected_logits_batch1 = Expectations(
            {
                ("cuda", None): [-1.5361, -1.6328, -1.5283, -1.5225, -1.5293, -1.5244, -1.5322, -1.5332, -1.5332, -1.5293, -1.5254, -1.5352, -1.5273, -1.5381, -1.5361],
                ("xpu", 3): [-1.5342, -1.6318, -1.5264, -1.5205, -1.5273, -1.5225, -1.5303, -1.5312, -1.5312, -1.5273, -1.5234, -1.5332, -1.5254, -1.5361, -1.5342],
            }
        )
        # fmt: on

        torch.testing.assert_close(
            logits[0, -1, :15],
            torch.tensor(expected_logits_batch0.get_expectation(), device=torch_device),
            rtol=1e-3,
            atol=1e-3,
        )
        torch.testing.assert_close(
            logits[1, -1, :15],
            torch.tensor(expected_logits_batch1.get_expectation(), device=torch_device),
            rtol=1e-3,
            atol=1e-3,
        )

    def test_model_generation(self):
        tokenizer = AutoTokenizer.from_pretrained("inceptionai/Jais-2-8B-Chat")
        model = Jais2ForCausalLM.from_pretrained(
            "inceptionai/Jais-2-8B-Chat", torch_dtype=torch.float16, device_map="auto"
        )
        input_text = "Simply put, the theory of relativity states that"
        model_inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
        model_inputs.pop("token_type_ids", None)

        generated_ids = model.generate(**model_inputs, max_new_tokens=32, do_sample=False)

Review comment (performance):
just a nit but i prefer keeping them like
```suggestion
        EXPECTED_LOGITS_BATCH0 = Expectations(
            {
                ("cuda", None): [-0.9751, -1.0918, -0.9600, -0.9526, -0.9600, -0.9551, -0.9624, -0.9644, -0.9644, -0.9600, -0.9561, -0.9658, -0.9585, -0.9688, -0.9663],
                ("xpu", 3): [-0.9692, -1.0859, -0.9541, -0.9468, -0.9546, -0.9492, -0.9570, -0.9585, -0.9585, -0.9541, -0.9507, -0.9604, -0.9526, -0.9634, -0.9609],
            }
        ).get_expectation()
        EXPECTED_LOGITS_BATCH1 = Expectations(
            {
                ("cuda", None): [-1.5361, -1.6328, -1.5283, -1.5225, -1.5293, -1.5244, -1.5322, -1.5332, -1.5332, -1.5293, -1.5254, -1.5352, -1.5273, -1.5381, -1.5361],
                ("xpu", 3): [-1.5342, -1.6318, -1.5264, -1.5205, -1.5273, -1.5225, -1.5303, -1.5312, -1.5312, -1.5273, -1.5234, -1.5332, -1.5254, -1.5361, -1.5342],
            }
        ).get_expectation()
```

Code after:

    @require_deterministic_for_xpu
    def test_model_logits(self):
        model_id = "inceptionai/Jais-2-8B-Chat"
        dummy_input = torch.LongTensor([[0, 0, 0, 0, 0, 0, 1, 2, 3], [1, 1, 2, 3, 4, 5, 6, 7, 8]]).to(torch_device)
        attention_mask = dummy_input.ne(0).to(torch.long)

        model = Jais2ForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")

        with torch.no_grad():
            logits = model(dummy_input, attention_mask=attention_mask).logits
        logits = logits.float()

        # fmt: off
        EXPECTED_LOGITS_BATCH0 = Expectations(
            {
                ("cuda", None): [-0.9751, -1.0918, -0.9600, -0.9526, -0.9600, -0.9551, -0.9624, -0.9644, -0.9644, -0.9600, -0.9561, -0.9658, -0.9585, -0.9688, -0.9663],
                ("xpu", 3): [-0.9692, -1.0859, -0.9541, -0.9468, -0.9546, -0.9492, -0.9570, -0.9585, -0.9585, -0.9541, -0.9507, -0.9604, -0.9526, -0.9634, -0.9609],
            }
        ).get_expectation()
        EXPECTED_LOGITS_BATCH1 = Expectations(
            {
                ("cuda", None): [-1.5361, -1.6328, -1.5283, -1.5225, -1.5293, -1.5244, -1.5322, -1.5332, -1.5332, -1.5293, -1.5254, -1.5352, -1.5273, -1.5381, -1.5361],
                ("xpu", 3): [-1.5342, -1.6318, -1.5264, -1.5205, -1.5273, -1.5225, -1.5303, -1.5312, -1.5312, -1.5273, -1.5234, -1.5332, -1.5254, -1.5361, -1.5342],
            }
        ).get_expectation()
        # fmt: on

        torch.testing.assert_close(
            logits[0, -1, :15],
            torch.tensor(EXPECTED_LOGITS_BATCH0, device=torch_device),
            rtol=1e-3,
            atol=1e-3,
        )
        torch.testing.assert_close(
            logits[1, -1, :15],
            torch.tensor(EXPECTED_LOGITS_BATCH1, device=torch_device),
            rtol=1e-3,
            atol=1e-3,
        )

    def test_model_generation(self):
        tokenizer = AutoTokenizer.from_pretrained("inceptionai/Jais-2-8B-Chat")
        model = Jais2ForCausalLM.from_pretrained(
            "inceptionai/Jais-2-8B-Chat", torch_dtype=torch.float16, device_map="auto"
        )
        input_text = "Simply put, the theory of relativity states that"
        model_inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
        model_inputs.pop("token_type_ids", None)

        generated_ids = model.generate(**model_inputs, max_new_tokens=32, do_sample=False)