fix models loading /saving issue

This commit is contained in:
Dobromir Popov
2025-09-02 16:05:44 +03:00
parent 1b54438082
commit 15cc694669
13 changed files with 2264 additions and 72 deletions

View File

@@ -18,8 +18,9 @@ class ModelRegistry:
self.models: Dict[str, ModelInterface] = {}
self.model_performance: Dict[str, Dict[str, Any]] = {}
def register_model(self, name: str, model: ModelInterface):
def register_model(self, model: ModelInterface):
"""Register a model in the registry"""
name = model.name
self.models[name] = model
self.model_performance[name] = {
'correct': 0,
@@ -28,6 +29,7 @@ class ModelRegistry:
'last_used': None
}
logger.info(f"Registered model: {name}")
return True
def get_model(self, name: str) -> Optional[ModelInterface]:
"""Get a model by name"""
@@ -65,6 +67,15 @@ class ModelRegistry:
return best_model
def unregister_model(self, name: str) -> bool:
"""Unregister a model from the registry"""
if name in self.models:
del self.models[name]
if name in self.model_performance:
del self.model_performance[name]
logger.info(f"Unregistered model: {name}")
return True
# Global model registry instance
_model_registry = ModelRegistry()
@@ -72,9 +83,9 @@ def get_model_registry() -> ModelRegistry:
"""Get the global model registry instance"""
return _model_registry
def register_model(name: str, model: ModelInterface):
def register_model(model: ModelInterface):
"""Register a model in the global registry"""
_model_registry.register_model(name, model)
return _model_registry.register_model(model)
def get_model(name: str) -> Optional[ModelInterface]:
"""Get a model from the global registry"""