def main()

in scripts/collect_models.py [0:0]

• 21 lines of code
• 1 McCabe index (conditional complexity)

def main() -> None:
    """
    Parses command line arguments and initiates the collection of model operators.
    """
    parser = argparse.ArgumentParser(
        description="Collect operators used in ONNX models from the Hugging Face Hub."
    )
    parser.add_argument("--model_limit", type=int, default=None, help="Maximum number of models to query from the Hub.")
    parser.add_argument("--limit", type=int, default=None, help="Maximum number of unique models to process.")
    parser.add_argument(
        "--from_cache", action="store_true",
        help="Only use local cache for loading models."
    )
    parser.add_argument(
        "--all_models", action="store_true",
        help="Include models from the transformers.js library in addition to tiny-random models."
    )
    args = parser.parse_args()

    collect_model_ops(
        model_limit=args.model_limit,
        from_cache=args.from_cache,
        limit=args.limit,
        include_all_models=args.all_models,
    )