@@ -141,14 +141,13 @@ def __init__(self, *args, **kwargs):
141
141
self .TINY_ONNX_SEQ2SEQ_MODEL_ID = "fxmarty/sshleifer-tiny-mbart-onnx"
142
142
self .TINY_ONNX_STABLE_DIFFUSION_MODEL_ID = "optimum-internal-testing/tiny-stable-diffusion-onnx"
143
143
144
-
145
144
def test_ella (self ):
146
- model_id = "optimum-internal-testing/tiny-random-llama"
147
- file_name = "model_optimized.onnx"
145
+ model_id = "optimum-internal-testing/tiny-random-llama"
146
+ file_name = "model_optimized.onnx"
148
147
149
148
model = ORTModelForCausalLM .from_pretrained (model_id , revision = "onnx" )
150
149
self .assertEqual (model .model_path .name , "model.onnx" )
151
-
150
+
152
151
model = ORTModelForCausalLM .from_pretrained (model_id , revision = "onnx" , file_name = file_name )
153
152
self .assertEqual (model .model_path .name , file_name )
154
153
@@ -161,13 +160,15 @@ def test_ella(self):
161
160
model = ORTModelForCausalLM .from_pretrained (model_id , revision = "merged-onnx" , subfolder = "subfolder" )
162
161
self .assertEqual (model .model_path .name , "model.onnx" )
163
162
164
- model = ORTModelForCausalLM .from_pretrained (model_id , revision = "merged-onnx" , subfolder = "subfolder" , file_name = file_name )
163
+ model = ORTModelForCausalLM .from_pretrained (
164
+ model_id , revision = "merged-onnx" , subfolder = "subfolder" , file_name = file_name
165
+ )
165
166
self .assertEqual (model .model_path .name , file_name )
166
-
167
- model = ORTModelForCausalLM .from_pretrained (model_id , revision = "merged-onnx" , file_name = "decoder_with_past_model.onnx" )
168
- self .assertEqual (model .model_path .name , "decoder_with_past_model.onnx" )
169
-
170
167
168
+ model = ORTModelForCausalLM .from_pretrained (
169
+ model_id , revision = "merged-onnx" , file_name = "decoder_with_past_model.onnx"
170
+ )
171
+ self .assertEqual (model .model_path .name , "decoder_with_past_model.onnx" )
171
172
172
173
def test_load_model_from_local_path (self ):
173
174
model = ORTModel .from_pretrained (self .LOCAL_MODEL_PATH )
0 commit comments