Merge pull request #108 from opendatahub-io/fix-convert-script

utils: prevent convert.py from erroring out
opendatahub-io · Oct 18, 2023 · 747af55 · 747af55
2 parents 46d41ab + a0203d2
commit 747af55
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/utils/convert.py b/utils/convert.py
@@ -3,12 +3,18 @@
 import caikit_nlp
 import argparse
 
-parser = argparse.ArgumentParser(prog="convert.py")
-parser.add_argument("--model-path", help="Path of the base HuggingFace model", )
-parser.add_argument("--model-save-path", help="Path to save the Caikit format model to")
 
-args = parser.parse_args()
+def main():
+    parser = argparse.ArgumentParser(prog="convert.py")
+    parser.add_argument("--model-path", help="Path of the base HuggingFace model", )
+    parser.add_argument("--model-save-path", help="Path to save the Caikit format model to")
 
-model = caikit_nlp.text_generation.TextGeneration.bootstrap(args.model_path)
+    args = parser.parse_args()
 
-model.save(model_path=args.model_save_path)
+    model = caikit_nlp.text_generation.TextGeneration.bootstrap(args.model_path)
+
+    model.save(model_path=args.model_save_path)
+
+
+if __name__ == "__main__":
+    main()