fix models loading /saving issue
This commit is contained in:
17
models.py
17
models.py
@@ -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"""
|
||||
|
Reference in New Issue
Block a user