Fixed issue of No custom_IO file found during compile through CLI

Author: abhishek-singh591

QEfficient/cloud/custom_yaml.py

@@ -1,5 +1,6 @@
-from pathlib import Path
 import warnings
+from pathlib import Path
+
 
 class CustomIOGenerator:
     """
@@ -86,7 +87,7 @@ def generate(self) -> dict:
         custom_io_lang = {}
         for name in output_names.get("lang", []):
             if name.endswith("_RetainedState"):
-                base = name[:-len("_RetainedState")]
+                base = name[: -len("_RetainedState")]
                 dtype = "float16" if "vision_embeds" in name else self.kv_cache_dtype
                 custom_io_lang[base] = dtype
                 custom_io_lang[name] = dtype
@@ -112,7 +113,7 @@ def generate(self) -> dict:
         custom_io = {}
         for name in output_names:
             if name.endswith("_RetainedState"):
-                base = name[:-len("_RetainedState")]
+                base = name[: -len("_RetainedState")]
                 dtype = "float16" if "pixel_values" in name else self.kv_cache_dtype
                 custom_io[base] = dtype
                 custom_io[name] = dtype
@@ -136,7 +137,7 @@ def generate(self) -> dict:
         custom_io = {"input_features": self.kv_cache_dtype}
         for name in output_names:
             if name.endswith("_RetainedState"):
-                base = name[:-len("_RetainedState")]
+                base = name[: -len("_RetainedState")]
                 custom_io[base] = self.kv_cache_dtype
                 custom_io[name] = self.kv_cache_dtype
         self.dump(custom_io, self.dtype_suffix)

QEfficient/cloud/export.py

@@ -12,7 +12,9 @@
 from QEfficient.base.common import QEFFCommonLoader
 from QEfficient.utils import check_and_assign_cache_dir
 from QEfficient.utils.logging_utils import logger
+
 from .custom_yaml import generate_custom_io
+
 # Specifically for Docker images.
 ROOT_DIR = os.path.dirname(os.path.abspath(""))
 
@@ -108,4 +110,4 @@ def main(
         help="Set full batch size to enable continuous batching mode, default is None",
     )
     args = parser.parse_args()
-    main(**args.__dict__)
+    main(**args.__dict__)

QEfficient/compile/compile_helper.py

@@ -171,7 +171,7 @@ def compile(
     Returns:
         :str: Path to compiled ``qpc`` package.
     """
-    
+
     if full_batch_size and batch_size != 1:
         raise ValueError("Only either batch_size or full_batch_size should be greater than one")
 
@@ -185,7 +185,7 @@ def compile(
         path=specialization_json_path,
         full_batch_size=full_batch_size,
     )
-    
+
     dtype_suffix = "int8" if mxint8 else "fp16"
     source_path = f"./custom_io_{dtype_suffix}.yaml"
     destination_path = os.path.join(os.path.dirname(qpc_path), f"custom_io_{dtype_suffix}.yaml")
@@ -244,4 +244,4 @@ def compile(
         else:
             logger.info(f"Compiled QPC files can be found here: {qpc_path}")
 
-    return qpc_path
+    return qpc_path