37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
|
from .registry import load_model as _load_model
|
||
|
|
from .types import UnifiedModel, GenerateConfig
|
||
|
|
|
||
|
|
# Announce the loaded model to the global spy so other parts of the app can access it
|
||
|
|
def load_model(*args, **kwargs):
|
||
|
|
"""
|
||
|
|
Proxy to the real registry.load_model that also announces the loaded model via __spy.
|
||
|
|
Accepts arbitrary args/kwargs to remain compatible with all loaders.
|
||
|
|
"""
|
||
|
|
model = _load_model(*args, **kwargs)
|
||
|
|
try:
|
||
|
|
import __spy as spy # local import to avoid hard dependency during tooling
|
||
|
|
# Best-effort extraction of a model name from common argument patterns
|
||
|
|
model_name = (
|
||
|
|
kwargs.get("source")
|
||
|
|
or kwargs.get("model")
|
||
|
|
or (args[0] if args else None)
|
||
|
|
or getattr(model, "name", None)
|
||
|
|
or "unknown"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Shallow capture of load parameters (omit non-serializable)
|
||
|
|
safe_params = {}
|
||
|
|
for k, v in kwargs.items():
|
||
|
|
try:
|
||
|
|
repr(v) # ensure it is representable
|
||
|
|
safe_params[k] = v
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
|
||
|
|
spy.set_model(str(model_name), model, **safe_params)
|
||
|
|
except Exception:
|
||
|
|
# Never let announcing break model loading
|
||
|
|
pass
|
||
|
|
return model
|
||
|
|
|
||
|
|
__all__ = ["load_model", "UnifiedModel", "GenerateConfig"]
|