revert causal lm tests

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
huggingface · Dec 5, 2024 · 9d0d9f2 · 9d0d9f2
1 parent 95ddb81
commit 9d0d9f2
Showing 1 changed file with 11 additions and 17 deletions.
diff --git a/tests/ipex/test_modeling.py b/tests/ipex/test_modeling.py
@@ -236,43 +236,37 @@ def test_compare_to_transformers(self, model_arch):
         dtype = torch.float16 if IS_XPU_AVAILABLE else torch.float32
         # Test model forward do not need cache.
         ipex_model = IPEXModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype)
-        device = ipex_model.device
         self.assertIsInstance(ipex_model.config, PretrainedConfig)
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         tokens = tokenizer(
             "This is a sample",
             return_tensors="pt",
             return_token_type_ids=False if model_arch in ("llama", "llama2") else None,
-        ).to(device)
-        outputs = ipex_model.generate(**tokens, max_new_tokens=1, return_dict_in_generate=True, output_logits=True)
+        )
+        inputs = ipex_model.prepare_inputs_for_generation(**tokens)
+        outputs = ipex_model(**inputs)
 
-        self.assertIsInstance(outputs.logits[0], torch.Tensor)
+        self.assertIsInstance(outputs.logits, torch.Tensor)
 
-        transformers_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype).to(device)
+        transformers_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype)
         with torch.no_grad():
-            transformers_outputs = transformers_model.generate(
-                **tokens, max_new_tokens=1, return_dict_in_generate=True, output_logits=True
-            )
+            transformers_outputs = transformers_model(**tokens)
 
         # Test re-load model
         with tempfile.TemporaryDirectory() as tmpdirname:
             ipex_model.save_pretrained(tmpdirname)
             loaded_model = self.IPEX_MODEL_CLASS.from_pretrained(tmpdirname, torch_dtype=dtype)
-            loaded_model_outputs = loaded_model.generate(
-                **tokens, max_new_tokens=1, return_dict_in_generate=True, output_logits=True
-            )
+            loaded_model_outputs = loaded_model(**inputs)
 
         # Test init method
         init_model = self.IPEX_MODEL_CLASS(transformers_model)
-        init_model_outputs = init_model.generate(
-            **tokens, max_new_tokens=1, return_dict_in_generate=True, output_logits=True
-        )
+        init_model_outputs = init_model(**inputs)
 
         # Compare tensor outputs
-        self.assertTrue(torch.allclose(outputs.logits[0], transformers_outputs.logits[0], atol=1e-4))
+        self.assertTrue(torch.allclose(outputs.logits, transformers_outputs.logits, atol=1e-4))
         # To avoid float pointing error
-        self.assertTrue(torch.allclose(outputs.logits[0], loaded_model_outputs.logits[0], atol=1e-7))
-        self.assertTrue(torch.allclose(outputs.logits[0], init_model_outputs.logits[0], atol=1e-7))
+        self.assertTrue(torch.allclose(outputs.logits, loaded_model_outputs.logits, atol=1e-7))
+        self.assertTrue(torch.allclose(outputs.logits, init_model_outputs.logits, atol=1e-7))
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     def test_pipeline(self, model_arch):