import argparse

from worlddisasterlm.optimization.export_onnx import export_onnx


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export WorldDisasterLM-8B to ONNX")
    parser.add_argument("--model-path", default="checkpoints/worlddisasterlm-8b-qlora")
    parser.add_argument("--output-path", default="artifacts/worlddisasterlm-8b.onnx")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    export_onnx(model_path=args.model_path, output_path=args.output_path)


if __name__ == "__main__":
    main()