mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 09:21:18 +08:00
fix mixed quant option (#1326)
This commit is contained in:
parent
595f5da146
commit
d2e02b3aae
@ -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,
|
||||||
)
|
)
|
||||||
|
Loading…
Reference in New Issue
Block a user