xTuring/src/xturing/models/mistral.py at 794f17cdadd618f86fd8e41e499bcb0c0cb0a3ab · stochasticai/xTuring · GitHub

1
2
3
4
5
6
7
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
from typing import Optional

from xturing.engines.mistral_engine import (
    Ministral314BEngine,
    Ministral314BInt8Engine,
    Ministral314BLoraEngine,
    Ministral314BLoraInt8Engine,
    Ministral314BLoraKbitEngine,
    Mistral7BEngine,
)
from xturing.models.causal import (
    CausalInt8Model,
    CausalLoraInt8Model,
    CausalLoraKbitModel,
    CausalLoraModel,
    CausalModel,
)


class Mistral7B(CausalModel):
    config_name: str = "mistral_7b"

    def __init__(self, weights_path: Optional[str] = None):
        super().__init__(Mistral7BEngine.config_name, weights_path)


class Ministral314B(CausalModel):
    config_name: str = "ministral_3_14b"

    def __init__(self, weights_path: Optional[str] = None):
        super().__init__(Ministral314BEngine.config_name, weights_path)


class Ministral314BLora(CausalLoraModel):
    config_name: str = "ministral_3_14b_lora"

    def __init__(self, weights_path: Optional[str] = None):
        super().__init__(Ministral314BLoraEngine.config_name, weights_path)


class Ministral314BInt8(CausalInt8Model):
    config_name: str = "ministral_3_14b_int8"

    def __init__(self, weights_path: Optional[str] = None):
        super().__init__(Ministral314BInt8Engine.config_name, weights_path)


class Ministral314BLoraInt8(CausalLoraInt8Model):
    config_name: str = "ministral_3_14b_lora_int8"

    def __init__(self, weights_path: Optional[str] = None):
        super().__init__(Ministral314BLoraInt8Engine.config_name, weights_path)


class Ministral314BLoraKbit(CausalLoraKbitModel):
    config_name: str = "ministral_3_14b_lora_kbit"

    def __init__(self, weights_path: Optional[str] = None):
        super().__init__(Ministral314BLoraKbitEngine.config_name, weights_path)