Skip to content

HFLMConfig

dendron.configs.hflm_config.HFLMConfig dataclass

Source code in src/dendron/configs/hflm_config.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@dataclass
class HFLMConfig:
    _model: Union[str, PreTrainedModel]
    _device: str
    _parallelize: bool
    _dtype: Union[str, torch.dtype]
    _load_in_4bit: bool
    _load_in_8bit: bool
    _add_bos_token: bool
    _offload_folder: Union[str, os.PathLike]

    def __init__(self, 
                 model: Union[str, PreTrainedModel], 
                 device: str = "cuda",
                 parallelize: bool = False, 
                 dtype: Union[str, torch.dtype] = torch.float16, 
                 load_in_4bit: bool = False, 
                 load_in_8bit: bool = False, 
                 add_bos_token: bool = False, 
                 offload_folder: Union[str, os.PathLike] = None):
        self._model = model
        self._device = device
        self._parallelize = parallelize
        self._dtype = dtype
        self._load_in_4bit = load_in_4bit
        self._load_in_8bit = load_in_8bit
        self._add_bos_token = add_bos_token
        self._offload_folder = offload_folder

    @property
    def model_name(self) -> str:
        if isinstance(self._model, str):
            return self._model
        elif isinstance(self._model, PreTrainedModel):
            return self._model.name_or_path

    @property
    def model(self) -> Union[str, PreTrainedModel]:
        return self._model

    @property
    def device(self) -> str:
        return self._device

    @property
    def parallelize(self) -> bool:
        return self._parallelize

    @property
    def dtype(self) -> Union[str, torch.dtype]:
        return self._dtype

    @property
    def load_in_4bit(self) -> bool:
        return self._load_in_4bit

    @property
    def load_in_8bit(self) -> bool:
        return self._load_in_8bit

    @property
    def add_bos_token(self) -> bool:
        return self._add_bos_token

    @property
    def offload_folder(self) -> Union[str, os.PathLike]:
        return self._offload_folder

    def to_dict(self) -> dict:
        return {
            "model_name": self.model_name,
            "device": self.device,
            "parallelize": self.parallelize,
            "dtype": self.dtype,
            "load_in_4bit": self.load_in_4bit,
            "load_in_8bit": self.load_in_8bit,
            "add_bos_token": self.add_bos_token,
            "offload_folder": self.offload_folder
        }

model property

device property

parallelize property

dtype property

load_in_4bit property

load_in_8bit property

add_bos_token property

to_dict()

Source code in src/dendron/configs/hflm_config.py 
76
77
78
79
80
81
82
83
84
85
86
def to_dict(self) -> dict:
    return {
        "model_name": self.model_name,
        "device": self.device,
        "parallelize": self.parallelize,
        "dtype": self.dtype,
        "load_in_4bit": self.load_in_4bit,
        "load_in_8bit": self.load_in_8bit,
        "add_bos_token": self.add_bos_token,
        "offload_folder": self.offload_folder
    }