|
24 | 24 |
|
25 | 25 | from paddlenlp.datasets import load_dataset
|
26 | 26 | from paddlenlp.experimental.autonlp import AutoTrainerForTextClassification
|
| 27 | +from paddlenlp.utils.env import PADDLE_INFERENCE_MODEL_SUFFIX |
27 | 28 | from tests.testing_utils import get_tests_dir, slow
|
28 | 29 |
|
29 | 30 | finetune_model_candidate = {
|
@@ -175,13 +176,13 @@ def test_multiclass(self, custom_model_candidate, hp_overrides):
|
175 | 176 | # test export
|
176 | 177 | temp_export_path = os.path.join(temp_dir_path, "test_export")
|
177 | 178 | auto_trainer.export(export_path=temp_export_path)
|
178 |
| - self.assertTrue(os.path.exists(os.path.join(temp_export_path, "model.pdmodel"))) |
| 179 | + self.assertTrue(os.path.exists(os.path.join(temp_export_path, f"model{PADDLE_INFERENCE_MODEL_SUFFIX}"))) |
179 | 180 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "taskflow_config.json")))
|
180 | 181 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "tokenizer_config.json")))
|
181 | 182 |
|
182 | 183 | # test export compress model
|
183 | 184 | auto_trainer.export(export_path=temp_export_path, compress=True)
|
184 |
| - self.assertTrue(os.path.exists(os.path.join(temp_export_path, "model.pdmodel"))) |
| 185 | + self.assertTrue(os.path.exists(os.path.join(temp_export_path, f"model{PADDLE_INFERENCE_MODEL_SUFFIX}"))) |
185 | 186 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "taskflow_config.json")))
|
186 | 187 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "tokenizer_config.json")))
|
187 | 188 |
|
@@ -400,13 +401,13 @@ def test_default_model_candidate(self, language, hp_overrides):
|
400 | 401 | # test export
|
401 | 402 | temp_export_path = os.path.join(temp_dir_path, "test_export")
|
402 | 403 | auto_trainer.export(export_path=temp_export_path)
|
403 |
| - self.assertTrue(os.path.exists(os.path.join(temp_export_path, "model.pdmodel"))) |
| 404 | + self.assertTrue(os.path.exists(os.path.join(temp_export_path, f"model{PADDLE_INFERENCE_MODEL_SUFFIX}"))) |
404 | 405 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "taskflow_config.json")))
|
405 | 406 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "tokenizer_config.json")))
|
406 | 407 |
|
407 | 408 | # test export compress model
|
408 | 409 | auto_trainer.export(export_path=temp_export_path, compress=True)
|
409 |
| - self.assertTrue(os.path.exists(os.path.join(temp_export_path, "model.pdmodel"))) |
| 410 | + self.assertTrue(os.path.exists(os.path.join(temp_export_path, f"model{PADDLE_INFERENCE_MODEL_SUFFIX}"))) |
410 | 411 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "taskflow_config.json")))
|
411 | 412 | self.assertTrue(os.path.exists(os.path.join(temp_export_path, "tokenizer_config.json")))
|
412 | 413 |
|
|
0 commit comments