adding quantization and model reusing for openvino backend

Author: Mohamed-Ashraf273

keras_hub/src/models/causal_lm.py

@@ -133,8 +133,13 @@ def make_generate_function(self):
 
         self.generate_function = self.generate_step
         if keras.config.backend() == "openvino":
+            import os
+            import tempfile
+
             import openvino as ov
             import openvino.runtime.opset14 as ov_opset
+            from nncf import CompressWeightsMode
+            from nncf import compress_weights
 
             from keras_hub.src.utils.openvino_utils import get_outputs
             from keras_hub.src.utils.openvino_utils import get_struct_outputs
@@ -143,17 +148,45 @@ def ov_infer(inputs, stop_token_ids, fn):
                 struct_params, struct_outputs = get_struct_outputs(
                     inputs, stop_token_ids, fn
                 )
-                parameters = [
-                    p.output.get_node() for p in tree.flatten(struct_params)
-                ]
-                results = [
-                    ov_opset.result(r.output)
-                    for r in tree.flatten(struct_outputs)
-                ]
-                core = ov.Core()
-                ov_model = ov.Model(results=results, parameters=parameters)
-                compile_ov_model = core.compile_model(ov_model, "CPU")
-                return get_outputs(inputs, struct_outputs, compile_ov_model)
+                if not hasattr(ov_infer, "compiled_model"):
+                    ov_infer.compiled_model = None
+                    parameters = [
+                        p.output.get_node() for p in tree.flatten(struct_params)
+                    ]
+                    results = [
+                        ov_opset.result(r.output)
+                        for r in tree.flatten(struct_outputs)
+                    ]
+                    core = ov.Core()
+                    ov_model = ov.Model(results=results, parameters=parameters)
+                    for ov_input in ov_model.inputs:
+                        rank = ov_input.get_partial_shape().rank.get_length()
+                        ov_input.get_node().set_partial_shape(
+                            ov.PartialShape([-1] * rank)
+                        )
+                    ov_model.validate_nodes_and_infer_types()
+                    with tempfile.TemporaryDirectory() as tmpdir:
+                        path = os.path.join(tmpdir, "ov_model.xml")
+                        ov.serialize(ov_model, path)
+                        del ov_model
+                        ov_model = core.read_model(path)
+                        group_sizes = [128, 64, 16, 4]
+                        for group_size in group_sizes:
+                            try:
+                                final_model = compress_weights(
+                                    ov_model,
+                                    mode=CompressWeightsMode.INT4_SYM,
+                                    group_size=group_size,
+                                )
+                                break
+                            except Exception:
+                                continue
+                    ov_infer.compile_ov_model = core.compile_model(
+                        final_model, "CPU"
+                    )
+                return get_outputs(
+                    inputs, struct_outputs, ov_infer.compile_ov_model
+                )
 
             def wrapped_generate_function(inputs, stop_token_ids=None):
                 inputs = tree.map_structure(ops.array, inputs)

keras_hub/src/utils/openvino_utils.py

@@ -30,21 +30,26 @@ def unpack_singleton(x):
     return x
 
 
-def parameterize_inputs(inputs):
+def parameterize_inputs(inputs, prefix=""):
     if isinstance(inputs, (list, tuple)):
-        return [parameterize_inputs(e) for e in inputs]
+        return [
+            parameterize_inputs(e, f"{prefix}{i}") for i, e in enumerate(inputs)
+        ]
     elif isinstance(inputs, dict):
-        return {k: parameterize_inputs(v) for k, v in inputs.items()}
+        return {k: parameterize_inputs(v, k) for k, v in inputs.items()}
     elif isinstance(inputs, np.ndarray):
         ov_type = OPENVINO_DTYPES[str(inputs.dtype)]
         ov_shape = list(inputs.shape)
         param = ov_opset.parameter(shape=ov_shape, dtype=ov_type)
+        param.set_friendly_name(prefix)
         return ops.convert_to_tensor(param.output(0))
     elif isinstance(inputs, (int, np.integer)):
         param = ov_opset.parameter(shape=[], dtype=ov.Type.i32)
+        param.set_friendly_name(prefix)
         return ops.convert_to_tensor(param.output(0))
     elif isinstance(inputs, (float, np.floating)):
         param = ov_opset.parameter(shape=[], dtype=ov.Type.f32)
+        param.set_friendly_name(prefix)
         return ops.convert_to_tensor(param.output(0))
     else:
         raise TypeError(f"Unknown input type: {type(inputs)}")

requirements-common.txt

@@ -19,4 +19,5 @@ tensorflow-datasets
 safetensors
 pillow
 openvino
+nncf
 transformers