dream

1
from typing import Optional, Union
2
from pydantic import BaseSettings, validator
3
from pathlib import Path
4

5
PRETRAINED_MODEL_PATH_FOLDER = ["/data/comet_commonsense", "pretrained_models"]
6

7

8
class AppConfig(BaseSettings):
9
    SERVICE_NAME: str
10
    SERVICE_PORT: int
11
    SENTRY_DSN: Optional[str]
12

13
    GRAPH: str
14
    PRETRAINED_MODEL: str
15
    DECODING_ALGO: str
16

17
    CUDA_VISIBLE_DEVICES: Union[int, str]
18

19
    @validator("PRETRAINED_MODEL")
20
    def create_full_model_path(cls, pretrained_model_pkl):
21
        model_file = Path(pretrained_model_pkl).name
22
        return Path(*PRETRAINED_MODEL_PATH_FOLDER, model_file)
23

24
    @validator("CUDA_VISIBLE_DEVICES")
25
    def device_validator(cls, device):
26
        if isinstance(device, int):
27
            return device
28
        return "cpu"
29

30
    @property
31
    def device(self):
32
        return "cuda:0" if self.CUDA_VISIBLE_DEVICES != "cpu" else "cpu"
33

34

35
settings = AppConfig()
36