optimizing the 'reusing model' logic

Author: Mohamed-Ashraf273

keras_hub/src/models/causal_lm.py

@@ -144,24 +144,9 @@ def ov_infer(inputs, stop_token_ids, fn):
                     inputs, stop_token_ids, fn
                 )
 
-                if not hasattr(ov_infer, "_compiled_models"):
-                    ov_infer._compiled_models = {}
-
-                # Create hash based on inputs, inputs shapes, and input dtypes
-                inputs_shapes = []
-                inputs_dtypes = []
-                for k, v in inputs.items():
-                    inputs_shapes.append(str(v.shape))
-                    inputs_dtypes.append(str(v.dtype))
-                model_signature = (
-                    f"inputs_{len(inputs)}_"
-                    f"shapes_{'_'.join(inputs_shapes)}_"
-                    f"dtypes_{'_'.join(inputs_dtypes)}_"
-                )
-
-                model_hash = hash(model_signature)
+                if not hasattr(ov_infer, "_compiled_model"):
+                    ov_infer._compiled_model = None
 
-                if model_hash not in ov_infer._compiled_models:
                     parameters = [
                         p.output.get_node() for p in tree.flatten(struct_params)
                     ]
@@ -171,11 +156,9 @@ def ov_infer(inputs, stop_token_ids, fn):
                     ]
 
                     ov_model = ov.Model(results=results, parameters=parameters)
-                    ov_infer._compiled_models[model_hash] = ov.compile_model(
-                        ov_model, "CPU"
-                    )
+                    ov_infer._compiled_model = ov.compile_model(ov_model, "CPU")
 
-                compile_ov_model = ov_infer._compiled_models[model_hash]
+                compile_ov_model = ov_infer._compiled_model
                 return get_outputs(inputs, struct_outputs, compile_ov_model)
 
             def wrapped_generate_function(inputs, stop_token_ids=None):

keras_hub/src/utils/openvino_utils.py

@@ -37,7 +37,7 @@ def parameterize_inputs(inputs):
         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)
+        ov_shape = [-1] * len(inputs.shape)
         param = ov_opset.parameter(shape=ov_shape, dtype=ov_type)
         return ops.convert_to_tensor(param.output(0))
     elif isinstance(inputs, (int, np.integer)):
@@ -50,9 +50,19 @@ def parameterize_inputs(inputs):
         raise TypeError(f"Unknown input type: {type(inputs)}")
 
 
+def reshape_params(params, inputs):
+    if isinstance(params, (list, tuple)):
+        return [reshape_params(p, i) for p, i in zip(params, inputs)]
+    elif isinstance(params, dict):
+        return {k: reshape_params(params[k], inputs[k]) for k in params}
+    else:
+        return ops.reshape(params, inputs.shape)
+
+
 def get_struct_outputs(inputs, stop_token_ids, fn):
     struct_params = parameterize_inputs(inputs)
-    struct_outputs = fn(struct_params, stop_token_ids)
+    struct_params_reshaped = reshape_params(struct_params, inputs)
+    struct_outputs = fn(struct_params_reshaped, stop_token_ids)
     return struct_params, struct_outputs