add suggested updates

Author: Mohamed-Ashraf273

conftest.py

@@ -1,12 +1,8 @@
 import os
-from pathlib import Path
 
 import keras
 import pytest
 
-from keras_hub.src.utils.openvino_utils import get_openvino_skip_reason
-from keras_hub.src.utils.openvino_utils import setup_openvino_test_config
-
 
 def pytest_addoption(parser):
     parser.addoption(
@@ -33,16 +29,27 @@ def pytest_addoption(parser):
         default=False,
         help="fail if a gpu is not present",
     )
-    parser.addoption(
-        "--auto_skip_training",
-        action="store_true",
-        default=True,
-        help="automatically skip tests with "
-        "training methods on non-trainable backends",
-    )
 
 
 def pytest_configure(config):
+    # Monkey-patch training methods for OpenVINO backend
+    if keras.config.backend() == "openvino":
+        # Store original methods in case we need to restore them
+        if not hasattr(keras.Model, "_original_compile"):
+            keras.Model._original_compile = keras.Model.compile
+            keras.Model._original_fit = keras.Model.fit
+            keras.Model._original_train_on_batch = keras.Model.train_on_batch
+
+        keras.Model.compile = lambda *args, **kwargs: pytest.skip(
+            "Model.compile() not supported on OpenVINO backend"
+        )
+        keras.Model.fit = lambda *args, **kwargs: pytest.skip(
+            "Model.fit() not supported on OpenVINO backend"
+        )
+        keras.Model.train_on_batch = lambda *args, **kwargs: pytest.skip(
+            "Model.train_on_batch() not supported on OpenVINO backend"
+        )
+
     # Verify that device has GPU and detected by backend
     if config.getoption("--check_gpu"):
         found_gpu = False
@@ -84,12 +91,9 @@ def pytest_configure(config):
 
 
 def pytest_collection_modifyitems(config, items):
-    openvino_supported_paths = None
-
     run_extra_large_tests = config.getoption("--run_extra_large")
     # Run large tests for --run_extra_large or --run_large.
     run_large_tests = config.getoption("--run_large") or run_extra_large_tests
-    auto_skip_training = config.getoption("--auto_skip_training")
 
     # Messages to annotate skipped tests with.
     skip_large = pytest.mark.skipif(
@@ -124,21 +128,58 @@ def pytest_collection_modifyitems(config, items):
         if "kaggle_key_required" in item.keywords:
             item.add_marker(kaggle_key_required)
 
-        # OpenVINO-specific skipping logic - whitelist-based approach
+        # OpenVINO-specific test skipping
         if keras.config.backend() == "openvino":
-            # OpenVINO backend configuration
-            if openvino_supported_paths is None:
-                openvino_supported_paths = setup_openvino_test_config(
-                    str(Path(__file__).parent)
+            test_name = item.name.split("[")[0]
+            test_path = str(item.fspath)
+
+            # OpenVINO supported test paths
+            openvino_supported_paths = [
+                "keras-hub/integration_tests",
+                "keras_hub/src/models/gemma",
+                "keras_hub/src/models/gpt2",
+                "keras_hub/src/models/mistral",
+                "keras_hub/src/samplers/serialization_test.py",
+                "keras_hub/src/tests/doc_tests/docstring_test.py",
+                "keras_hub/src/tokenizers",
+                "keras_hub/src/utils",
+            ]
+
+            # Skip specific problematic test methods
+            specific_skipping_tests = {
+                "test_backbone_basics": "Requires trainable backend",
+                "test_score_loss": "Non-implemented roll operation",
+                "test_layer_behaviors": "Requires trainable backend",
+            }
+
+            if test_name in specific_skipping_tests:
+                item.add_marker(
+                    pytest.mark.skipif(
+                        True,
+                        reason="OpenVINO: "
+                        f"{specific_skipping_tests[test_name]}",
+                    )
                 )
-            skip_reason = get_openvino_skip_reason(
-                item,
-                openvino_supported_paths,
-                auto_skip_training,
+                continue
+
+            parts = test_path.replace("\\", "/").split("/")
+            try:
+                keras_hub_idx = parts.index("keras_hub")
+                relative_test_path = "/".join(parts[keras_hub_idx:])
+            except ValueError:
+                relative_test_path = test_path
+
+            is_whitelisted = any(
+                relative_test_path == supported_path
+                or relative_test_path.startswith(supported_path + "/")
+                for supported_path in openvino_supported_paths
             )
-            if skip_reason:
+
+            if not is_whitelisted:
                 item.add_marker(
-                    pytest.mark.skipif(True, reason=f"OpenVINO: {skip_reason}")
+                    pytest.mark.skipif(
+                        True, reason="OpenVINO: File/directory not in whitelist"
+                    )
                 )
 
 

keras_hub/src/samplers/beam_sampler.py

@@ -95,15 +95,15 @@ def unflatten_beams(x):
         )
         log_probs = flatten_beams(ops.repeat(log_probs, batch_size, axis=0))
 
-        def cond(prompt, cache, index, log_probs):
+        def cond(prompt, cache, index, mask, log_probs):
             if stop_token_ids is None:
                 return True
             # Stop if all sequences have produced a *new* stop token.
             end_tokens = any_equal(prompt, stop_token_ids, ~mask)
             prompt_done = ops.any(end_tokens, axis=-1)
             return ops.logical_not(ops.all(prompt_done))
 
-        def body(prompt, cache, index, log_probs):
+        def body(prompt, cache, index, mask, log_probs):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
             vocab_size = ops.shape(logits)[-1]
@@ -150,12 +150,12 @@ def gather_beams(x):
             next_token = next_token[:, None]
             prompt = ops.slice_update(prompt, [0, index], next_token)
             # Return the iteration of the loop state.
-            return (prompt, cache, index + 1, log_probs)
+            return (prompt, cache, index + 1, mask, log_probs)
 
-        prompt, _, _, log_probs = self.run_loop(
+        prompt, _, _, _, log_probs = self.run_loop(
             cond=cond,
             body=body,
-            loop_vars=(prompt, cache, index, log_probs),
+            loop_vars=(prompt, cache, index, mask, log_probs),
             maximum_iterations=(max_length - index),
             model=model,
         )