File size: 701 Bytes
d68e65a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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