16
16
17
17
import torch
18
18
from parameterized import parameterized
19
-
20
- # TODO : add more tasks
21
19
from transformers import (
22
20
AutoModelForCausalLM ,
23
21
AutoModelForQuestionAnswering ,
26
24
AutoTokenizer ,
27
25
pipeline ,
28
26
)
27
+ from utils_tests import MODEL_NAMES
29
28
30
29
from optimum .intel import inference_mode as ipex_inference_mode
31
30
from optimum .intel .ipex .modeling_base import IPEXModel
32
31
33
32
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
-
50
33
_CLASSIFICATION_TASK_TO_AUTOMODELS = {
51
34
"text-classification" : AutoModelForSequenceClassification ,
52
35
"token-classification" : AutoModelForTokenClassification ,
53
36
}
54
37
55
38
56
- class IPEXIntegrationTest (unittest .TestCase ):
57
- CLASSIFICATION_SUPPORTED_ARCHITECTURES = (
39
+ class IPEXClassificationTest (unittest .TestCase ):
40
+ SUPPORTED_ARCHITECTURES = (
58
41
"bert" ,
59
42
"distilbert" ,
60
43
"roberta" ,
61
44
)
62
45
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 )
74
54
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 = (
76
65
"bert" ,
77
66
"distilbert" ,
78
67
"roberta" ,
79
68
)
80
69
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 ):
83
72
model_id = MODEL_NAMES [model_arch ]
84
73
tokenizer = AutoTokenizer .from_pretrained (model_id )
85
74
model = AutoModelForQuestionAnswering .from_pretrained (model_id , torch_dtype = torch .float32 )
@@ -95,24 +84,22 @@ def test_question_answering_pipeline_inference(self, model_arch):
95
84
self .assertEqual (outputs ["start" ], outputs_ipex ["start" ])
96
85
self .assertEqual (outputs ["end" ], outputs_ipex ["end" ])
97
86
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 )
106
87
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
+ )
113
100
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 ):
116
103
model_id = MODEL_NAMES [model_arch ]
117
104
model = AutoModelForCausalLM .from_pretrained (model_id , torch_dtype = torch .float32 , return_dict = False )
118
105
model = model .eval ()
0 commit comments