cognisafe-backend / app /ml_models /classifier_path_loader.py
zyriean's picture
add app
d68e65a verified
from pathlib import Path
path_to_model = {"toxic-bert": Path.cwd() / "app" / "ml_models" / "toxic-bert"}
class ClassifierPathLoader:
def __init__(self) -> None:
self.model_name = None
self.model_path = None
def set_model(self, model_name: str) -> None:
if model_name not in path_to_model:
raise KeyError(f"Model '{model_name}' not found in path registry.")
self.model_name = model_name
self.model_path = Path(path_to_model[model_name])
def get_model_path(self) -> Path:
if self.model_path is None:
raise RuntimeError("Model not set. Call `set_model()` first.")
return self.model_path