cleaning causal_lm and avoid keras.src imports

Author: Mohamed-Ashraf273

keras_hub/src/models/causal_lm.py

@@ -133,55 +133,11 @@ def make_generate_function(self):
 
         self.generate_function = self.generate_step
         if keras.config.backend() == "openvino":
-            import numpy as np
             import openvino as ov
             import openvino.runtime.opset14 as ov_opset
-            from keras.src.backend.openvino.core import OPENVINO_DTYPES
-            from keras.src.backend.openvino.core import OpenVINOKerasTensor
-
-            def unpack_singleton(x):
-                if isinstance(x, (list, tuple)) and len(x) == 1:
-                    return x[0]
-                return x
-
-            def parameterize_inputs(inputs):
-                if isinstance(inputs, (list, tuple)):
-                    return [parameterize_inputs(e) for e in inputs]
-                elif isinstance(inputs, dict):
-                    return {
-                        k: parameterize_inputs(v) 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)
-                    return OpenVINOKerasTensor(param.output(0))
-                elif isinstance(inputs, (int, np.integer)):
-                    param = ov_opset.parameter(shape=[], dtype=ov.Type.i32)
-                    return OpenVINOKerasTensor(param.output(0))
-                elif isinstance(inputs, (float, np.floating)):
-                    param = ov_opset.parameter(shape=[], dtype=ov.Type.f32)
-                    return OpenVINOKerasTensor(param.output(0))
-                else:
-                    raise TypeError(f"Unknown input type: {type(inputs)}")
-
-            def get_struct_outputs(inputs, stop_token_ids, fn):
-                struct_params = parameterize_inputs(inputs)
-                struct_outputs = fn(struct_params, stop_token_ids)
-                return struct_params, struct_outputs
-
-            def get_outputs_from_model(
-                inputs, struct_outputs, compile_ov_model
-            ):
-                flatten_inputs = tree.flatten(inputs)
-                assert OpenVINOKerasTensor not in inputs, (
-                    "inputs should be numpy arrays"
-                )
-                outputs = compile_ov_model(flatten_inputs)
-                outputs = unpack_singleton(
-                    tree.pack_sequence_as(struct_outputs, outputs.to_tuple())
-                )
-                return outputs
+
+            from keras_hub.src.utils.openvino_utils import get_outputs
+            from keras_hub.src.utils.openvino_utils import get_struct_outputs
 
             def ov_infer(inputs, stop_token_ids, fn):
                 struct_params, struct_outputs = get_struct_outputs(
@@ -220,14 +176,11 @@ def ov_infer(inputs, stop_token_ids, fn):
                     )
 
                 compile_ov_model = ov_infer._compiled_models[model_hash]
-                return get_outputs_from_model(
-                    inputs, struct_outputs, compile_ov_model
-                )
+                return get_outputs(inputs, struct_outputs, compile_ov_model)
 
             def wrapped_generate_function(inputs, stop_token_ids=None):
                 for k, v in inputs.items():
-                    if isinstance(v, OpenVINOKerasTensor):
-                        inputs[k] = ops.array(v)
+                    inputs[k] = ops.array(v)
                 return ov_infer(inputs, stop_token_ids, self.generate_step)
 
             self.generate_function = wrapped_generate_function

keras_hub/src/utils/openvino_utils.py

@@ -0,0 +1,66 @@
+import numpy as np
+import openvino as ov
+import openvino.runtime.opset14 as ov_opset
+from keras import tree
+from keras.src.backend.openvino.core import OpenVINOKerasTensor
+
+OPENVINO_DTYPES = {
+    "float16": ov.Type.f16,
+    "float32": ov.Type.f32,
+    "float64": ov.Type.f64,
+    "uint8": ov.Type.u8,
+    "uint16": ov.Type.u16,
+    "uint32": ov.Type.u32,
+    "uint64": ov.Type.u64,
+    "int8": ov.Type.i8,
+    "int16": ov.Type.i16,
+    "int32": ov.Type.i32,
+    "int64": ov.Type.i64,
+    "bfloat16": ov.Type.bf16,
+    "bool": ov.Type.boolean,
+    "float8_e4m3fn": ov.Type.f8e4m3,
+    "float8_e5m2": ov.Type.f8e5m2,
+    "string": ov.Type.string,
+}
+
+
+def unpack_singleton(x):
+    if isinstance(x, (list, tuple)) and len(x) == 1:
+        return x[0]
+    return x
+
+
+def parameterize_inputs(inputs):
+    if isinstance(inputs, (list, tuple)):
+        return [parameterize_inputs(e) for e in inputs]
+    elif isinstance(inputs, dict):
+        return {k: parameterize_inputs(v) 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)
+        return OpenVINOKerasTensor(param.output(0))
+    elif isinstance(inputs, (int, np.integer)):
+        param = ov_opset.parameter(shape=[], dtype=ov.Type.i32)
+        return OpenVINOKerasTensor(param.output(0))
+    elif isinstance(inputs, (float, np.floating)):
+        param = ov_opset.parameter(shape=[], dtype=ov.Type.f32)
+        return OpenVINOKerasTensor(param.output(0))
+    else:
+        raise TypeError(f"Unknown input type: {type(inputs)}")
+
+
+def get_struct_outputs(inputs, stop_token_ids, fn):
+    struct_params = parameterize_inputs(inputs)
+    struct_outputs = fn(struct_params, stop_token_ids)
+    return struct_params, struct_outputs
+
+
+def get_outputs(inputs, struct_outputs, compile_ov_model):
+    flatten_inputs = tree.flatten(inputs)
+    assert OpenVINOKerasTensor not in inputs, "inputs should be numpy arrays"
+    outputs = compile_ov_model(flatten_inputs)
+    outputs = unpack_singleton(
+        tree.pack_sequence_as(struct_outputs, outputs.to_tuple())
+    )
+    return outputs