From 2de42ed179774f4e161438dd7b64165caabf007d Mon Sep 17 00:00:00 2001 From: Your Name Date: Thu, 27 Jun 2024 21:30:36 +0000 Subject: [PATCH] Renamed the new function and added a new one to delete the model and the model data --- .../python/transformers/test_parity_mixtral_moe.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py b/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py index 00a928f4faf85..848607ae45771 100644 --- a/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py +++ b/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py @@ -42,13 +42,16 @@ def print_tensor(name, numpy_array): print(f"const std::vector {name} = {value_string_of(numpy_array)};") -def create_onnx_graph(model, model_path): +def save_model_to_disk(model, model_path): external_data_path = "mixtral_moe.onnx" + ".data" onnx.save_model( model, model_path, save_as_external_data=True, all_tensors_to_one_file=True, location=external_data_path ) - return model_path + +def delete_model_data(external_data): + os.remove("mixtral_moe.onnx") + os.remove(external_data) def create_moe_onnx_graph( @@ -139,9 +142,9 @@ def create_moe_onnx_graph( model = helper.make_model(graph) model_path = "mixtral_moe.onnx" - save_model = create_onnx_graph(model, model_path) + save_model_to_disk(model, model_path) - return save_model + return model_path class ClassInstantier(OrderedDict): @@ -434,9 +437,8 @@ def test_mixtral_moe_benchmark(self): mixtral_moe = MixtralSparseMoeBlock(config, batch_size, sequence_length) mixtral_moe.benchmark() - os.remove("mixtral_moe.onnx") external_data_path = "mixtral_moe.onnx" + ".data" - os.remove(external_data_path) + delete_model_data(external_data_path) if __name__ == "__main__":