1616
1717import torch
1818from parameterized import parameterized
19-
20- # TODO : add more tasks
2119from transformers import (
2220 AutoModelForCausalLM ,
2321 AutoModelForQuestionAnswering ,
2624 AutoTokenizer ,
2725 pipeline ,
2826)
27+ from utils_tests import MODEL_NAMES
2928
3029from optimum .intel import inference_mode as ipex_inference_mode
3130from optimum .intel .ipex .modeling_base import IPEXModel
3231
3332
34- MODEL_NAMES = {
35- "bert" : "hf-internal-testing/tiny-random-bert" ,
36- "bloom" : "hf-internal-testing/tiny-random-BloomModel" ,
37- "distilbert" : "hf-internal-testing/tiny-random-distilbert" ,
38- "roberta" : "hf-internal-testing/tiny-random-roberta" ,
39- "gptj" : "hf-internal-testing/tiny-random-gptj" ,
40- "gpt2" : "hf-internal-testing/tiny-random-gpt2" ,
41- "gpt_neo" : "hf-internal-testing/tiny-random-GPTNeoModel" ,
42- "gpt_neox" : "hf-internal-testing/tiny-random-GPTNeoXForCausalLM" ,
43- "gpt_bigcode" : "hf-internal-testing/tiny-random-GPTBigCodeModel" ,
44- "llama" : "fxmarty/tiny-llama-fast-tokenizer" ,
45- "llama2" : "Jiqing/tiny_random_llama2" ,
46- "opt" : "hf-internal-testing/tiny-random-OPTModel" ,
47- "mpt" : "hf-internal-testing/tiny-random-MptForCausalLM" ,
48- }
49-
5033_CLASSIFICATION_TASK_TO_AUTOMODELS = {
5134 "text-classification" : AutoModelForSequenceClassification ,
5235 "token-classification" : AutoModelForTokenClassification ,
5336}
5437
5538
56- class IPEXIntegrationTest (unittest .TestCase ):
57- CLASSIFICATION_SUPPORTED_ARCHITECTURES = (
39+ class IPEXClassificationTest (unittest .TestCase ):
40+ SUPPORTED_ARCHITECTURES = (
5841 "bert" ,
5942 "distilbert" ,
6043 "roberta" ,
6144 )
6245
63- TEXT_GENERATION_SUPPORTED_ARCHITECTURES = (
64- "bloom" ,
65- "gptj" ,
66- "gpt2" ,
67- "gpt_neo" ,
68- "gpt_bigcode" ,
69- "llama" ,
70- "llama2" ,
71- "opt" ,
72- "mpt" ,
73- )
46+ @parameterized .expand (SUPPORTED_ARCHITECTURES )
47+ def test_pipeline_inference (self , model_arch ):
48+ model_id = MODEL_NAMES [model_arch ]
49+ tokenizer = AutoTokenizer .from_pretrained (model_id )
50+ inputs = "This is a sample input"
51+ for task , auto_model_class in _CLASSIFICATION_TASK_TO_AUTOMODELS .items ():
52+ model = auto_model_class .from_pretrained (model_id , torch_dtype = torch .float32 )
53+ pipe = pipeline (task , model = model , tokenizer = tokenizer )
7454
75- QA_SUPPORTED_ARCHITECTURES = (
55+ with torch .inference_mode ():
56+ outputs = pipe (inputs )
57+ with ipex_inference_mode (pipe , dtype = model .config .torch_dtype , verbose = False , jit = True ) as ipex_pipe :
58+ outputs_ipex = ipex_pipe (inputs )
59+ self .assertTrue (isinstance (ipex_pipe .model ._optimized .model , torch .jit .RecursiveScriptModule ))
60+ self .assertEqual (outputs [0 ]["score" ], outputs_ipex [0 ]["score" ])
61+
62+
63+ class IPEXQuestionAnsweringTest (unittest .TestCase ):
64+ SUPPORTED_ARCHITECTURES = (
7665 "bert" ,
7766 "distilbert" ,
7867 "roberta" ,
7968 )
8069
81- @parameterized .expand (QA_SUPPORTED_ARCHITECTURES )
82- def test_question_answering_pipeline_inference (self , model_arch ):
70+ @parameterized .expand (SUPPORTED_ARCHITECTURES )
71+ def test_pipeline_inference (self , model_arch ):
8372 model_id = MODEL_NAMES [model_arch ]
8473 tokenizer = AutoTokenizer .from_pretrained (model_id )
8574 model = AutoModelForQuestionAnswering .from_pretrained (model_id , torch_dtype = torch .float32 )
@@ -95,24 +84,22 @@ def test_question_answering_pipeline_inference(self, model_arch):
9584 self .assertEqual (outputs ["start" ], outputs_ipex ["start" ])
9685 self .assertEqual (outputs ["end" ], outputs_ipex ["end" ])
9786
98- @parameterized .expand (CLASSIFICATION_SUPPORTED_ARCHITECTURES )
99- def test_classification_pipeline_inference (self , model_arch ):
100- model_id = MODEL_NAMES [model_arch ]
101- tokenizer = AutoTokenizer .from_pretrained (model_id )
102- inputs = "This is a sample input"
103- for task , auto_model_class in _CLASSIFICATION_TASK_TO_AUTOMODELS .items ():
104- model = auto_model_class .from_pretrained (model_id , torch_dtype = torch .float32 )
105- pipe = pipeline (task , model = model , tokenizer = tokenizer )
10687
107- with torch .inference_mode ():
108- outputs = pipe (inputs )
109- with ipex_inference_mode (pipe , dtype = model .config .torch_dtype , verbose = False , jit = True ) as ipex_pipe :
110- outputs_ipex = ipex_pipe (inputs )
111- self .assertTrue (isinstance (ipex_pipe .model ._optimized .model , torch .jit .RecursiveScriptModule ))
112- self .assertEqual (outputs [0 ]["score" ], outputs_ipex [0 ]["score" ])
88+ class IPEXTextGenerationTest (unittest .TestCase ):
89+ SUPPORTED_ARCHITECTURES = (
90+ "bloom" ,
91+ "gptj" ,
92+ "gpt2" ,
93+ "gpt_neo" ,
94+ "gpt_bigcode" ,
95+ "llama" ,
96+ "llama2" ,
97+ "opt" ,
98+ "mpt" ,
99+ )
113100
114- @parameterized .expand (TEXT_GENERATION_SUPPORTED_ARCHITECTURES )
115- def test_text_generation_pipeline_inference (self , model_arch ):
101+ @parameterized .expand (SUPPORTED_ARCHITECTURES )
102+ def test_pipeline_inference (self , model_arch ):
116103 model_id = MODEL_NAMES [model_arch ]
117104 model = AutoModelForCausalLM .from_pretrained (model_id , torch_dtype = torch .float32 , return_dict = False )
118105 model = model .eval ()
0 commit comments