update PR

Author: Mohamed-Ashraf273

keras_hub/src/models/gemma/gemma_causal_lm.py

@@ -258,9 +258,6 @@ def next(prompt, cache, index):
             cache_update_index = index - 1
             batch_size = ops.shape(prompt)[0]
             prompt = ops.slice(prompt, [0, cache_update_index], [batch_size, 1])
-            if keras.config.backend() == "openvino":
-                # Avoid returning dynamic shape by openvino slice
-                prompt = ops.reshape(prompt, [batch_size, 1])
             logits, hidden_states, cache = self.call_with_cache(
                 prompt,
                 cache,

keras_hub/src/models/gpt2/gpt2_causal_lm.py

@@ -246,9 +246,6 @@ def next(prompt, cache, index):
             cache_update_index = index - 1
             batch_size = ops.shape(prompt)[0]
             prompt = ops.slice(prompt, [0, cache_update_index], [batch_size, 1])
-            if keras.config.backend() == "openvino":
-                # Avoid returning dynamic shape by openvino slice
-                prompt = ops.reshape(prompt, [batch_size, 1])
             logits, hidden_states, cache = self.call_with_cache(
                 prompt,
                 cache,

keras_hub/src/models/mistral/mistral_causal_lm.py

@@ -145,9 +145,6 @@ def next(prompt, cache, index):
             cache_update_index = index - 1
             batch_size = ops.shape(prompt)[0]
             prompt = ops.slice(prompt, [0, cache_update_index], [batch_size, 1])
-            if keras.config.backend() == "openvino":
-                # Avoid returning dynamic shape by openvino slice
-                prompt = ops.reshape(prompt, [batch_size, 1])
             logits, hidden_states, cache = self.call_with_cache(
                 prompt,
                 cache,

keras_hub/src/utils/openvino_utils.py

@@ -2,7 +2,6 @@
 import functools
 from pathlib import Path
 
-import keras
 from keras import tree
 
 from keras_hub.src.utils.keras_utils import print_msg
@@ -44,10 +43,9 @@ def setup_openvino_test_config(config_file_path):
     Returns:
         list: Supported paths (whitelist) for OpenVINO testing.
     """
-    supported_paths = load_openvino_supported_tools(
+    return load_openvino_supported_tools(
         Path(config_file_path) / "openvino_supported_tests.txt"
     )
-    return supported_paths
 
 
 @functools.lru_cache(maxsize=256)
@@ -147,7 +145,7 @@ def should_auto_skip_training_test(item):
     Returns:
         bool: True if should skip, False otherwise.
     """
-    if not item.fspath.basename.endswith(".py"):
+    if not str(item.fspath).endswith(".py"):
         return False
     test_name = item.name.split("[")[0]
     return _contains_training_methods(str(item.fspath), test_name)
@@ -166,9 +164,6 @@ def get_openvino_skip_reason(item, supported_paths, auto_skip_training=True):
     Returns:
         str or None: Skip reason if should skip, None otherwise.
     """
-    if keras.config.backend() != "openvino":
-        return None
-
     test_name = item.name.split("[")[0]
     test_path = str(item.fspath)
 
@@ -187,35 +182,30 @@ def get_openvino_skip_reason(item, supported_paths, auto_skip_training=True):
 
     # Priority 3: Whitelist-based approach - only test supported paths
     if supported_paths:
-        # Check if this test file/directory is in the whitelist
-        # Convert test path to relative path format for comparison
-        test_path_parts = test_path.replace("\\", "/").split("/")
-        # Find keras_hub index and create relative path
+        parts = test_path.replace("\\", "/").split("/")
         try:
-            keras_hub_idx = test_path_parts.index("keras_hub")
-            relative_test_path = "/".join(test_path_parts[keras_hub_idx:])
+            keras_hub_idx = parts.index("keras_hub")
+            relative_test_path = "/".join(parts[keras_hub_idx:])
         except ValueError:
-            relative_test_path = test_path
+            relative_test_path = test_path  # fall back to absolute
 
         for supported_path in supported_paths:
-            # Exact match or directory prefix match
             if (
                 relative_test_path == supported_path
                 or relative_test_path.startswith(supported_path + "/")
-                or relative_test_path.startswith(supported_path + "\\")
             ):
-                return None  # Found in whitelist, don't skip
+                return None  # in whitelist
 
-        # Not found in whitelist, skip it
         return "File/directory not in OpenVINO whitelist"
+
     return None
 
 
 def get_device():
     """Detect and return the best available OpenVINO device.
 
     Returns:
-        tuple: (core, device) where device is "GPU" or "CPU".
+        str: "GPU" if available, otherwise "CPU".
     """
     return "GPU" if "GPU" in core.available_devices else "CPU"
 
@@ -232,16 +222,17 @@ def compile_model(struct_params, struct_outputs, device, model_dtype):
     Returns:
         Compiled OpenVINO model ready for inference.
     """
-    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)]
+    flat_params = tree.flatten(struct_params)
+    flat_outputs = tree.flatten(struct_outputs)
+    parameters = [p.output.get_node() for p in flat_params]
+    results = [ov_opset.result(r.output) for r in flat_outputs]
     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()
     config = {"INFERENCE_PRECISION_HINT": model_dtype}
-    compiled_model = core.compile_model(ov_model, device, config)
-    return compiled_model
+    return core.compile_model(ov_model, device, config)
 
 
 def get_outputs(inputs, struct_outputs, compiled_ov_model, unpack_singleton):
@@ -257,9 +248,9 @@ def get_outputs(inputs, struct_outputs, compiled_ov_model, unpack_singleton):
         Structured model outputs matching expected format.
     """
     flatten_inputs = tree.flatten(inputs)
-    outputs = compiled_ov_model(flatten_inputs).to_tuple()
-    outputs = unpack_singleton(tree.pack_sequence_as(struct_outputs, outputs))
-    return outputs
+    raw = compiled_ov_model(flatten_inputs).to_tuple()
+    packed = tree.pack_sequence_as(struct_outputs, raw)
+    return unpack_singleton(packed)
 
 
 def ov_infer(model, inputs, stop_token_ids, fn):
@@ -280,7 +271,7 @@ def ov_infer(model, inputs, stop_token_ids, fn):
     """
     device = get_device()
 
-    # Try to use existing compiled model
+    # Try to use existing compiled model for the same device
     if (
         getattr(model, "ov_compiled_model", None) is not None
         and getattr(model, "ov_device", None) is not None
@@ -298,8 +289,8 @@ def ov_infer(model, inputs, stop_token_ids, fn):
                 "WARNING: OpenVINO inference \033[1mFAILED\033[0m, "
                 "recompiling model and trying again.\n" + str(e)
             )
-            del model.ov_compiled_model
-            del model.struct_outputs
+            model.ov_compiled_model = None
+            model.struct_outputs = None
 
     # Compile a new model
     struct_params = model._parameterize_data(inputs)
@@ -309,7 +300,6 @@ def ov_infer(model, inputs, stop_token_ids, fn):
     model.ov_compiled_model = compile_model(
         struct_params, model.struct_outputs, device, model_dtype
     )
-
     return get_outputs(
         inputs,
         model.struct_outputs,