[SW-20808] Make sure save&load format is an Enum object (#58)

* [SW-20808] Make sure save&load format is an Enum object

Signed-off-by: Xin He <xinhe3@habana.ai>

* Update save_load_entry.py

---------

Signed-off-by: Xin He <xinhe3@habana.ai>
Co-authored-by: Xin He <xinhe3@habana.ai>
Signed-off-by: Xin He <xinhe3@habana.ai>

Author: xin3he

neural_compressor/torch/algorithms/fp8_quant/save_load.py

@@ -348,6 +348,10 @@ def save_for_multi_devices(model, checkpoint_dir="saved_results", format="huggin
         checkpoint_dir (str, optional): path to checkpoint. Defaults to "saved_results".
         format (str, optional): defaults to 'huggingface'.
     """
+    format = get_enum_from_format(format)
+    assert format == SaveLoadFormat.HUGGINGFACE, (
+        "Currently, only huggingface models are supported." + "Please set format='huggingface'."
+    )
     from safetensors.torch import save_file as safe_save_file
     if format == SaveLoadFormat.VLLM:
         import transformers

neural_compressor/torch/algorithms/weight_only/save_load.py

@@ -61,6 +61,7 @@ def save(model, output_dir="./saved_results", format=SaveLoadFormat.DEFAULT, **k
             - tokenizer (Tokenizer, optional): The tokenizer to be saved along with the model (only applicable for 'huggingface' format).
             - max_shard_size (str, optional): The maximum size for each shard (only applicable for 'huggingface' format). Defaults to "5GB".
     """
+    format = get_enum_from_format(format)
     os.makedirs(output_dir, exist_ok=True)
     cur_accelerator.synchronize()
     if format == SaveLoadFormat.HUGGINGFACE:  # pragma: no cover
@@ -128,6 +129,7 @@ def load(model_name_or_path, original_model=None, format=SaveLoadFormat.DEFAULT,
     Returns:
         torch.nn.Module: quantized model
     """
+    format = get_enum_from_format(format)
     model_loader = WOQModelLoader(model_name_or_path, original_model, format, device, **kwargs)
     model = model_loader.load_woq_model()
     return model

neural_compressor/torch/quantization/save_load_entry.py

@@ -26,7 +26,7 @@
     RTNConfig,
     TEQConfig,
 )
-from neural_compressor.torch.utils import SaveLoadFormat
+from neural_compressor.torch.utils import SaveLoadFormat, get_enum_from_format
 
 config_name_mapping = {
     FP8_QUANT: FP8Config,
@@ -45,6 +45,7 @@ def save(model, checkpoint_dir="saved_results", format="default"):
             quantized by llm-compressor(https://github.com/vllm-project/llm-compressor).
             Defaults to "default".
     """
+    format = get_enum_from_format(format)
     config_mapping = model.qconfig
     config_object = config_mapping[next(iter(config_mapping))]
     # fp8_quant
@@ -104,7 +105,8 @@ def load(model_name_or_path, original_model=None, format="default", device="cpu"
     Returns:
         The quantized model
     """
-    if format == SaveLoadFormat.DEFAULT.value:
+    format = get_enum_from_format(format)
+    if format == SaveLoadFormat.DEFAULT:
         from neural_compressor.common.base_config import ConfigRegistry
 
         qconfig_file_path = os.path.join(os.path.abspath(os.path.expanduser(model_name_or_path)), "qconfig.json")
@@ -133,7 +135,7 @@ def load(model_name_or_path, original_model=None, format="default", device="cpu"
                     model_name_or_path, original_model, format=SaveLoadFormat.DEFAULT, device=device
                 )
                 return qmodel.to(device)
-    elif format == SaveLoadFormat.HUGGINGFACE.value:
+    elif format == SaveLoadFormat.HUGGINGFACE:
         import transformers
 
         config = transformers.AutoConfig.from_pretrained(model_name_or_path, **kwargs)
@@ -156,4 +158,4 @@ def load(model_name_or_path, original_model=None, format="default", device="cpu"
             qmodel = weight_only.load(model_name_or_path, format=SaveLoadFormat.HUGGINGFACE, device=device, **kwargs)
             return qmodel.to(device)
     else:
-        raise ValueError("`format` in load function can only be 'huggingface' or 'default', but get {}".format(format))
+        assert False, "This code path should never be reached."

neural_compressor/torch/utils/utility.py

@@ -32,7 +32,7 @@
     detect_processor_type_based_on_hw,
     logger,
 )
-from neural_compressor.torch.utils import is_optimum_habana_available, is_transformers_imported
+from neural_compressor.torch.utils import is_optimum_habana_available, is_transformers_imported, SaveLoadFormat
 
 if is_transformers_imported():
     import transformers
@@ -711,3 +711,16 @@ def forward_wrapper(model, input):
     else:
         output = model(input)
     return output
+
+
+def get_enum_from_format(format):
+    """Make sure Save&Load format is an Enum object."""
+    if isinstance(format, SaveLoadFormat):
+        return format
+    for obj in SaveLoadFormat:
+        if format == obj.value:
+            return obj
+        elif format.upper() == obj.name:
+            return obj
+    raise ValueError(
+        f"Invalid format value ('{format}'). Enter one of [{[m.name for m in SaveLoadFormat]}]")