[djl-import] Fixes missing arguments for onnx import (#3458)

deepjavalibrary · Sep 7, 2024 · 61ffdbb · 61ffdbb
1 parent 584a4b1
commit 61ffdbb
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py
@@ -78,6 +78,16 @@ def importer_args():
                         "--cpu-only",
                         action='store_true',
                         help="Only validate jit traced model on CPU")
+    parser.add_argument("--optimize",
+                        choices=["O1", "O2", "O3", "O4"],
+                        help="Optimization option for ONNX models")
+    parser.add_argument(
+        "--device",
+        help='The device to use to do the export. Defaults to "cpu".')
+    parser.add_argument(
+        "--dtype",
+        choices=["fp32", "fp16", "bf16"],
+        help="The floating point precision to use for the export.")
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument(
         "-c",