AI tries to fix it's own tests v2

app/core/engine.py

@@ -6,7 +6,7 @@ import time
 from dataclasses import dataclass
 import inspect
 from threading import Lock, Thread
-from typing import Generator, Iterable, List, Optional, Sequence, Tuple
+from typing import Callable, Generator, Iterable, List, Optional, Sequence, Tuple
 
 # Heavy libraries (torch, transformers) are imported lazily in functions
 # to keep application startup snappy.
@@ -24,6 +24,32 @@ def _unwrap_bound_callable(callable_obj):
     return getattr(callable_obj, "__func__", callable_obj)
 
 
+def _install_tie_weights_compat_patch(model) -> Callable[[], None]:
+    """Patch model class tie_weights to ignore newer unexpected kwargs."""
+    if not hasattr(model, "tie_weights"):
+        return lambda: None
+
+    model_cls = model.__class__
+    had_local_attr = "tie_weights" in model_cls.__dict__
+    original_attr = model_cls.__dict__.get("tie_weights")
+    original_callable = _unwrap_bound_callable(getattr(model_cls, "tie_weights"))
+
+    def _compat_tie_weights(self, *args, **kwargs):
+        kwargs.pop("missing_keys", None)
+        kwargs.pop("recompute_mapping", None)
+        return original_callable(self, *args, **kwargs)
+
+    setattr(model_cls, "tie_weights", _compat_tie_weights)
+
+    def _restore() -> None:
+        if had_local_attr:
+            setattr(model_cls, "tie_weights", original_attr)
+        else:
+            delattr(model_cls, "tie_weights")
+
+    return _restore
+
+
 def _lazy_import_torch():  # pragma: no cover - indirection
     import torch
 
@@ -199,17 +225,12 @@ class _ModelHandle:
             _restore_loader_attr = _orig_load_pretrained_func
         
         def _patched_load_pretrained_func(cls, model, *args, **kwargs):
-            # Patch model.tie_weights to accept and ignore unexpected kwargs
-            orig_tie_weights = model.tie_weights
-            def _compat_tie_weights(*tw_args, **tw_kwargs):
-                tw_kwargs.pop("missing_keys", None)
-                tw_kwargs.pop("recompute_mapping", None)
-                return orig_tie_weights(*tw_args, **tw_kwargs)
-            model.tie_weights = _compat_tie_weights
+            # Patch tie_weights to absorb kwargs passed by newer transformers.
+            restore_tie_weights = _install_tie_weights_compat_patch(model)
             try:
                 return _orig_load_pretrained_func(cls, model, *args, **kwargs)
             finally:
-                model.tie_weights = orig_tie_weights
+                restore_tie_weights()
         
         modeling_utils.PreTrainedModel._load_pretrained_model = classmethod(_patched_load_pretrained_func)
         try:

tests/test_core_helpers.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 
 import types
 
+import pytest
+
 from app.core import engine, tokens
 from app.core.prompting import DEFAULT_SYSTEM_PROMPT, render_chat_prompt
 from app.schemas.chat import ChatMessage
@@ -110,3 +112,23 @@ def test_unwrap_bound_callable_extracts_underlying_method_function() -> None:
     unwrapped = engine._unwrap_bound_callable(Demo.loader)
     assert callable(unwrapped)
     assert unwrapped.__name__ == "loader"
+
+
+def test_install_tie_weights_compat_patch_strips_unexpected_kwargs() -> None:
+    class Demo:
+        def tie_weights(self):  # pragma: no cover - shape-only test
+            return "ok"
+
+    demo = Demo()
+
+    with pytest.raises(TypeError):
+        demo.tie_weights(missing_keys=set(), recompute_mapping=False)
+
+    restore = engine._install_tie_weights_compat_patch(demo)
+    try:
+        assert demo.tie_weights(missing_keys=set(), recompute_mapping=False) == "ok"
+    finally:
+        restore()
+
+    with pytest.raises(TypeError):
+        demo.tie_weights(missing_keys=set(), recompute_mapping=False)