Improve: Graphical display of ONNX conversion script logs

convert_bert_onnx.py

@@ -1,12 +1,16 @@
 # usage: .venv/bin/python convert_bert_onnx.py --language JP
 # ref: https://github.com/tuna2134/sbv2-api/blob/main/convert/convert_deberta.py
 
+import time
 from argparse import ArgumentParser
 from pathlib import Path
 
 import onnx
 import torch
-from onnxsim import simplify
+from onnxsim import model_info, simplify
+from rich import print
+from rich.rule import Rule
+from rich.style import Style
 from torch import nn
 from transformers.convert_slow_tokenizer import BertConverter
 
@@ -15,6 +19,7 @@ from style_bert_vits2.nlp import bert_models
 
 
 if __name__ == "__main__":
+    start_time = time.time()
     parser = ArgumentParser()
     parser.add_argument("--language", default=Languages.JP)
     args = parser.parse_args()
@@ -25,11 +30,18 @@ if __name__ == "__main__":
     onnx_temp_model_path = Path(pretrained_model_name_or_path) / f"model_temp.onnx"
     onnx_optimized_model_path = Path(pretrained_model_name_or_path) / f"model.onnx"
     tokenizer_json_path = Path(pretrained_model_name_or_path) / "tokenizer.json"
+    print(Rule(characters="=", style=Style(color="blue")))
+    print(f"[bold cyan]Language:[/bold cyan] {language.name}")
+    print(f"[bold cyan]Pretrained model:[/bold cyan] {pretrained_model_name_or_path}")
+    print(Rule(characters="=", style=Style(color="blue")))
 
     # トークナイザーを Fast Tokenizer 用形式に変換して保存
     tokenizer = bert_models.load_tokenizer(language)
     converter = BertConverter(tokenizer)
     converter.converted().save(str(tokenizer_json_path))
+    print(Rule(characters="=", style=Style(color="blue")))
+    print(f"[bold green]Tokenizer JSON saved to {tokenizer_json_path}[/bold green]")
+    print(Rule(characters="=", style=Style(color="blue")))
 
     # TODO: JP, ZH は変換できるが、EN は途中で強制終了されてしまい変換できない
     class ONNXBert(nn.Module):
@@ -52,6 +64,10 @@ if __name__ == "__main__":
     inputs = tokenizer("今日はいい天気ですね", return_tensors="pt")
 
     # モデルを ONNX に変換
+    print(Rule(characters="=", style=Style(color="blue")))
+    print(f"[bold cyan]Exporting ONNX model...[/bold cyan]")
+    print(Rule(characters="=", style=Style(color="blue")))
+    export_start_time = time.time()
     torch.onnx.export(
         model=model,
         args=(inputs["input_ids"], inputs["token_type_ids"], inputs["attention_mask"]),
@@ -64,12 +80,26 @@ if __name__ == "__main__":
             "attention_mask": {1: "batch_size"},
         },
     )
+    print(
+        f"[bold green]ONNX model exported to {onnx_temp_model_path} ({time.time() - export_start_time:.2f}s)[/bold green]"
+    )
 
     # ONNX モデルを最適化
+    print(Rule(characters="=", style=Style(color="blue")))
+    print(f"[bold cyan]Optimizing ONNX model...[/bold cyan]")
+    print(Rule(characters="=", style=Style(color="blue")))
+    optimize_start_time = time.time()
     onnx_model = onnx.load(onnx_temp_model_path)
     simplified_onnx_model, check = simplify(onnx_model)
     onnx.save(simplified_onnx_model, onnx_optimized_model_path)
-
-    # 最適化前の ONNX モデルを削除
     onnx_temp_model_path.unlink()
-    print(f"ONNX model optimized and saved to {onnx_optimized_model_path}")
+    print(
+        f"[bold green]ONNX model optimized and saved to {onnx_optimized_model_path} ({time.time() - optimize_start_time:.2f}s)[/bold green]"
+    )
+    print(
+        f"[bold]Total Time: {time.time() - start_time:.2f}s / Size: {onnx_optimized_model_path.stat().st_size / 1024 / 1024:.2f}MB[/bold]"
+    )
+    print(Rule(characters="=", style=Style(color="blue")))
+    print("[bold cyan]Optimized model info:[/bold cyan]")
+    model_info.print_simplifying_info(onnx_model, simplified_onnx_model)
+    print(Rule(characters="=", style=Style(color="blue")))