fix mixed quant option

This commit is contained in:
Awni Hannun 2025-03-07 06:55:35 -08:00
parent e7267d30f8
commit c72811da57

View File

@ -1,27 +1,23 @@
# Copyright © 2023-2024 Apple Inc. # Copyright © 2023-2024 Apple Inc.
import argparse import argparse
from enum import Enum
from .utils import convert, mixed_2_6, mixed_3_6 from . import utils
from .utils import convert
QUANT_RECIPES = [
class MixedQuants(Enum): "mixed_2_6",
mixed_3_6 = "mixed_3_6" "mixed_3_6",
mixed_2_6 = "mixed_2_6" ]
@classmethod
def recipe_names(cls):
return [member.name for member in cls]
def quant_args(arg): def quant_args(arg):
try: if arg not in QUANT_RECIPES:
return MixedQuants[arg].value
except KeyError:
raise argparse.ArgumentTypeError( raise argparse.ArgumentTypeError(
f"Invalid q-recipe {arg!r}. Choose from: {MixedQuants.recipe_names()}" f"Invalid q-recipe {arg!r}. Choose from: {QUANT_RECIPES}"
) )
else:
return getattr(utils, arg)
def configure_parser() -> argparse.ArgumentParser: def configure_parser() -> argparse.ArgumentParser:
@ -50,7 +46,7 @@ def configure_parser() -> argparse.ArgumentParser:
) )
parser.add_argument( parser.add_argument(
"--quant-predicate", "--quant-predicate",
help=f"Mixed-bit quantization recipe. Choices: {MixedQuants.recipe_names()}", help=f"Mixed-bit quantization recipe. Choices: {QUANT_RECIPES}",
type=quant_args, type=quant_args,
required=False, required=False,
) )