czczup commited on
Commit
362dca6
1 Parent(s): 01eb2ce

Upload folder using huggingface_hub

Browse files
config.json CHANGED
@@ -4,6 +4,10 @@
4
  "architectures": [
5
  "InternVLChatModel"
6
  ],
 
 
 
 
7
  "downsample_ratio": 0.5,
8
  "force_image_size": 448,
9
  "llm_config": {
 
4
  "architectures": [
5
  "InternVLChatModel"
6
  ],
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_internvl_chat.InternVLChatConfig",
9
+ "AutoModel": "modeling_internvl_chat.InternVLChatModel"
10
+ }
11
  "downsample_ratio": 0.5,
12
  "force_image_size": 448,
13
  "llm_config": {
configuration_intern_vit.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # --------------------------------------------------------
2
+ # InternVL
3
+ # Copyright (c) 2023 OpenGVLab
4
+ # Licensed under The MIT License [see LICENSE for details]
5
+ # --------------------------------------------------------
6
+ import os
7
+ from typing import Union
8
+
9
+ from transformers.configuration_utils import PretrainedConfig
10
+ from transformers.utils import logging
11
+
12
+ logger = logging.get_logger(__name__)
13
+
14
+
15
+ class InternVisionConfig(PretrainedConfig):
16
+ r"""
17
+ This is the configuration class to store the configuration of a [`InternVisionModel`]. It is used to
18
+ instantiate a vision encoder according to the specified arguments, defining the model architecture.
19
+
20
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
21
+ documentation from [`PretrainedConfig`] for more information.
22
+
23
+ Args:
24
+ num_channels (`int`, *optional*, defaults to 3):
25
+ Number of color channels in the input images (e.g., 3 for RGB).
26
+ patch_size (`int`, *optional*, defaults to 14):
27
+ The size (resolution) of each patch.
28
+ image_size (`int`, *optional*, defaults to 224):
29
+ The size (resolution) of each image.
30
+ qkv_bias (`bool`, *optional*, defaults to `False`):
31
+ Whether to add a bias to the queries and values in the self-attention layers.
32
+ hidden_size (`int`, *optional*, defaults to 3200):
33
+ Dimensionality of the encoder layers and the pooler layer.
34
+ num_attention_heads (`int`, *optional*, defaults to 25):
35
+ Number of attention heads for each attention layer in the Transformer encoder.
36
+ intermediate_size (`int`, *optional*, defaults to 12800):
37
+ Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
38
+ qk_normalization (`bool`, *optional*, defaults to `True`):
39
+ Whether to normalize the queries and keys in the self-attention layers.
40
+ num_hidden_layers (`int`, *optional*, defaults to 48):
41
+ Number of hidden layers in the Transformer encoder.
42
+ use_flash_attn (`bool`, *optional*, defaults to `True`):
43
+ Whether to use flash attention mechanism.
44
+ hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
45
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
46
+ `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported.
47
+ layer_norm_eps (`float`, *optional*, defaults to 1e-6):
48
+ The epsilon used by the layer normalization layers.
49
+ dropout (`float`, *optional*, defaults to 0.0):
50
+ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
51
+ drop_path_rate (`float`, *optional*, defaults to 0.0):
52
+ Dropout rate for stochastic depth.
53
+ attention_dropout (`float`, *optional*, defaults to 0.0):
54
+ The dropout ratio for the attention probabilities.
55
+ initializer_range (`float`, *optional*, defaults to 0.02):
56
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
57
+ initializer_factor (`float`, *optional*, defaults to 0.1):
58
+ A factor for layer scale.
59
+ """
60
+
61
+ model_type = 'intern_vit_6b'
62
+
63
+ def __init__(
64
+ self,
65
+ num_channels=3,
66
+ patch_size=14,
67
+ image_size=224,
68
+ qkv_bias=False,
69
+ hidden_size=3200,
70
+ num_attention_heads=25,
71
+ intermediate_size=12800,
72
+ qk_normalization=True,
73
+ num_hidden_layers=48,
74
+ use_flash_attn=True,
75
+ hidden_act='gelu',
76
+ layer_norm_eps=1e-6,
77
+ dropout=0.0,
78
+ drop_path_rate=0.0,
79
+ attention_dropout=0.0,
80
+ initializer_range=0.02,
81
+ initializer_factor=0.1,
82
+ **kwargs,
83
+ ):
84
+ super().__init__(**kwargs)
85
+
86
+ self.hidden_size = hidden_size
87
+ self.intermediate_size = intermediate_size
88
+ self.dropout = dropout
89
+ self.drop_path_rate = drop_path_rate
90
+ self.num_hidden_layers = num_hidden_layers
91
+ self.num_attention_heads = num_attention_heads
92
+ self.num_channels = num_channels
93
+ self.patch_size = patch_size
94
+ self.image_size = image_size
95
+ self.initializer_range = initializer_range
96
+ self.initializer_factor = initializer_factor
97
+ self.attention_dropout = attention_dropout
98
+ self.layer_norm_eps = layer_norm_eps
99
+ self.hidden_act = hidden_act
100
+ self.qkv_bias = qkv_bias
101
+ self.qk_normalization = qk_normalization
102
+ self.use_flash_attn = use_flash_attn
103
+
104
+ @classmethod
105
+ def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> 'PretrainedConfig':
106
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
107
+
108
+ if 'vision_config' in config_dict:
109
+ config_dict = config_dict['vision_config']
110
+
111
+ if 'model_type' in config_dict and hasattr(cls, 'model_type') and config_dict['model_type'] != cls.model_type:
112
+ logger.warning(
113
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
114
+ f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
115
+ )
116
+
117
+ return cls.from_dict(config_dict, **kwargs)
configuration_internvl_chat.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # --------------------------------------------------------
2
+ # InternVL
3
+ # Copyright (c) 2023 OpenGVLab
4
+ # Licensed under The MIT License [see LICENSE for details]
5
+ # --------------------------------------------------------
6
+
7
+ import copy
8
+
9
+ from transformers import LlamaConfig
10
+ from transformers.configuration_utils import PretrainedConfig
11
+ from transformers.utils import logging
12
+
13
+ from .configuration_intern_vit import InternVisionConfig
14
+
15
+ logger = logging.get_logger(__name__)
16
+
17
+
18
+ class InternVLChatConfig(PretrainedConfig):
19
+ model_type = 'internvl_chat'
20
+ is_composition = True
21
+
22
+ def __init__(
23
+ self,
24
+ vision_config=None,
25
+ llm_config=None,
26
+ use_backbone_lora=0,
27
+ use_llm_lora=0,
28
+ pad2square=False,
29
+ select_layer=-4,
30
+ force_image_size=None,
31
+ downsample_ratio=0.5,
32
+ template=None,
33
+ **kwargs):
34
+ super().__init__(**kwargs)
35
+
36
+ if vision_config is None:
37
+ vision_config = {}
38
+ logger.info('vision_config is None. Initializing the InternVisionConfig with default values.')
39
+
40
+ if llm_config is None:
41
+ llm_config = {}
42
+ logger.info('llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`).')
43
+
44
+ self.vision_config = InternVisionConfig(**vision_config)
45
+ self.llm_config = LlamaConfig(**llm_config)
46
+ self.use_backbone_lora = use_backbone_lora
47
+ self.use_llm_lora = use_llm_lora
48
+ self.pad2square = pad2square
49
+ self.select_layer = select_layer
50
+ self.force_image_size = force_image_size
51
+ self.downsample_ratio = downsample_ratio
52
+ self.template = template
53
+
54
+ def to_dict(self):
55
+ """
56
+ Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
57
+
58
+ Returns:
59
+ `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
60
+ """
61
+ output = copy.deepcopy(self.__dict__)
62
+ output['vision_config'] = self.vision_config.to_dict()
63
+ output['llm_config'] = self.llm_config.to_dict()
64
+ output['model_type'] = self.__class__.model_type
65
+ output['use_backbone_lora'] = self.use_backbone_lora
66
+ output['use_llm_lora'] = self.use_llm_lora
67
+ output['pad2square'] = self.pad2square
68
+ output['select_layer'] = self.select_layer
69
+ output['force_image_size'] = self.force_image_size
70
+ output['downsample_ratio'] = self.downsample_ratio
71
+ output['template'] = self.template
72
+
73
+ return output
conversation.py ADDED
@@ -0,0 +1,1256 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Conversation prompt templates.
3
+
4
+ We kindly request that you import fastchat instead of copying this file if you wish to use it.
5
+ If you have any changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
6
+ """
7
+
8
+ import dataclasses
9
+ from enum import IntEnum, auto
10
+ from typing import Any, Dict, List, Tuple, Union
11
+
12
+
13
+ class SeparatorStyle(IntEnum):
14
+ """Separator styles."""
15
+
16
+ ADD_COLON_SINGLE = auto()
17
+ ADD_COLON_TWO = auto()
18
+ ADD_COLON_SPACE_SINGLE = auto()
19
+ NO_COLON_SINGLE = auto()
20
+ NO_COLON_TWO = auto()
21
+ ADD_NEW_LINE_SINGLE = auto()
22
+ LLAMA2 = auto()
23
+ CHATGLM = auto()
24
+ CHATML = auto()
25
+ CHATINTERN = auto()
26
+ DOLLY = auto()
27
+ RWKV = auto()
28
+ PHOENIX = auto()
29
+ ROBIN = auto()
30
+ FALCON_CHAT = auto()
31
+ CHATGLM3 = auto()
32
+ INTERNVL_ZH = auto()
33
+
34
+
35
+ @dataclasses.dataclass
36
+ class Conversation:
37
+ """A class that manages prompt templates and keeps all conversation history."""
38
+
39
+ # The name of this template
40
+ name: str
41
+ # The template of the system prompt
42
+ system_template: str = '{system_message}'
43
+ # The system message
44
+ system_message: str = ''
45
+ # The names of two roles
46
+ roles: Tuple[str] = ('USER', 'ASSISTANT')
47
+ # All messages. Each item is (role, message).
48
+ messages: List[List[str]] = ()
49
+ # The number of few shot examples
50
+ offset: int = 0
51
+ # The separator style and configurations
52
+ sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
53
+ sep: str = '\n'
54
+ sep2: str = None
55
+ # Stop criteria (the default one is EOS token)
56
+ stop_str: Union[str, List[str]] = None
57
+ # Stops generation if meeting any token in this list
58
+ stop_token_ids: List[int] = None
59
+
60
+ def get_prompt(self) -> str:
61
+ """Get the prompt for generation."""
62
+ system_prompt = self.system_template.format(system_message=self.system_message)
63
+ if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
64
+ ret = system_prompt + self.sep
65
+ for role, message in self.messages:
66
+ if message:
67
+ ret += role + ': ' + message + self.sep
68
+ else:
69
+ ret += role + ':'
70
+ return ret
71
+ elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
72
+ seps = [self.sep, self.sep2]
73
+ ret = system_prompt + seps[0]
74
+ for i, (role, message) in enumerate(self.messages):
75
+ if message:
76
+ ret += role + ': ' + message + seps[i % 2]
77
+ else:
78
+ ret += role + ':'
79
+ return ret
80
+ elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
81
+ ret = system_prompt + self.sep
82
+ for role, message in self.messages:
83
+ if message:
84
+ ret += role + ': ' + message + self.sep
85
+ else:
86
+ ret += role + ': ' # must be end with a space
87
+ return ret
88
+ elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
89
+ ret = '' if system_prompt == '' else system_prompt + self.sep
90
+ for role, message in self.messages:
91
+ if message:
92
+ ret += role + '\n' + message + self.sep
93
+ else:
94
+ ret += role + '\n'
95
+ return ret
96
+ elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
97
+ ret = system_prompt
98
+ for role, message in self.messages:
99
+ if message:
100
+ ret += role + message + self.sep
101
+ else:
102
+ ret += role
103
+ return ret
104
+ elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
105
+ seps = [self.sep, self.sep2]
106
+ ret = system_prompt
107
+ for i, (role, message) in enumerate(self.messages):
108
+ if message:
109
+ ret += role + message + seps[i % 2]
110
+ else:
111
+ ret += role
112
+ return ret
113
+ elif self.sep_style == SeparatorStyle.RWKV:
114
+ ret = system_prompt
115
+ for i, (role, message) in enumerate(self.messages):
116
+ if message:
117
+ ret += (
118
+ role
119
+ + ': '
120
+ + message.replace('\r\n', '\n').replace('\n\n', '\n')
121
+ )
122
+ ret += '\n\n'
123
+ else:
124
+ ret += role + ':'
125
+ return ret
126
+ elif self.sep_style == SeparatorStyle.LLAMA2:
127
+ seps = [self.sep, self.sep2]
128
+ if self.system_message:
129
+ ret = system_prompt
130
+ else:
131
+ ret = '[INST] '
132
+ for i, (role, message) in enumerate(self.messages):
133
+ tag = self.roles[i % 2]
134
+ if message:
135
+ if i == 0:
136
+ ret += message + ' '
137
+ else:
138
+ ret += tag + ' ' + message + seps[i % 2]
139
+ else:
140
+ ret += tag
141
+ return ret
142
+ elif self.sep_style == SeparatorStyle.CHATGLM:
143
+ # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
144
+ # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
145
+ round_add_n = 1 if self.name == 'chatglm2' else 0
146
+ if system_prompt:
147
+ ret = system_prompt + self.sep
148
+ else:
149
+ ret = ''
150
+
151
+ for i, (role, message) in enumerate(self.messages):
152
+ if i % 2 == 0:
153
+ ret += f'[Round {i//2 + round_add_n}]{self.sep}'
154
+
155
+ if message:
156
+ ret += f'{role}:{message}{self.sep}'
157
+ else:
158
+ ret += f'{role}:'
159
+ return ret
160
+ elif self.sep_style == SeparatorStyle.CHATML:
161
+ ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
162
+ for role, message in self.messages:
163
+ if message:
164
+ ret += role + '\n' + message + self.sep + '\n'
165
+ else:
166
+ ret += role + '\n'
167
+ return ret
168
+ elif self.sep_style == SeparatorStyle.CHATGLM3:
169
+ ret = ''
170
+ if self.system_message:
171
+ ret += system_prompt
172
+ for role, message in self.messages:
173
+ if message:
174
+ ret += role + '\n' + ' ' + message
175
+ else:
176
+ ret += role
177
+ return ret
178
+ elif self.sep_style == SeparatorStyle.CHATINTERN:
179
+ # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
180
+ seps = [self.sep, self.sep2]
181
+ ret = system_prompt
182
+ for i, (role, message) in enumerate(self.messages):
183
+ # if i % 2 == 0:
184
+ # ret += "<s>"
185
+ if message:
186
+ ret += role + ':' + message + seps[i % 2] + '\n'
187
+ else:
188
+ ret += role + ':'
189
+ return ret
190
+ elif self.sep_style == SeparatorStyle.DOLLY:
191
+ seps = [self.sep, self.sep2]
192
+ ret = system_prompt
193
+ for i, (role, message) in enumerate(self.messages):
194
+ if message:
195
+ ret += role + ':\n' + message + seps[i % 2]
196
+ if i % 2 == 1:
197
+ ret += '\n\n'
198
+ else:
199
+ ret += role + ':\n'
200
+ return ret
201
+ elif self.sep_style == SeparatorStyle.PHOENIX:
202
+ ret = system_prompt
203
+ for role, message in self.messages:
204
+ if message:
205
+ ret += role + ': ' + '<s>' + message + '</s>'
206
+ else:
207
+ ret += role + ': ' + '<s>'
208
+ return ret
209
+ elif self.sep_style == SeparatorStyle.ROBIN:
210
+ ret = system_prompt + self.sep
211
+ for role, message in self.messages:
212
+ if message:
213
+ ret += role + ':\n' + message + self.sep
214
+ else:
215
+ ret += role + ':\n'
216
+ return ret
217
+ elif self.sep_style == SeparatorStyle.FALCON_CHAT:
218
+ ret = ''
219
+ if self.system_message:
220
+ ret += system_prompt + self.sep
221
+ for role, message in self.messages:
222
+ if message:
223
+ ret += role + ': ' + message + self.sep
224
+ else:
225
+ ret += role + ':'
226
+
227
+ return ret
228
+ elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
229
+ seps = [self.sep, self.sep2]
230
+ ret = self.system_message + seps[0]
231
+ for i, (role, message) in enumerate(self.messages):
232
+ if message:
233
+ ret += role + ': ' + message + seps[i % 2]
234
+ else:
235
+ ret += role + ':'
236
+ return ret
237
+ else:
238
+ raise ValueError(f'Invalid style: {self.sep_style}')
239
+
240
+ def set_system_message(self, system_message: str):
241
+ """Set the system message."""
242
+ self.system_message = system_message
243
+
244
+ def append_message(self, role: str, message: str):
245
+ """Append a new message."""
246
+ self.messages.append([role, message])
247
+
248
+ def update_last_message(self, message: str):
249
+ """Update the last output.
250
+
251
+ The last message is typically set to be None when constructing the prompt,
252
+ so we need to update it in-place after getting the response from a model.
253
+ """
254
+ self.messages[-1][1] = message
255
+
256
+ def to_gradio_chatbot(self):
257
+ """Convert the conversation to gradio chatbot format."""
258
+ ret = []
259
+ for i, (role, msg) in enumerate(self.messages[self.offset :]):
260
+ if i % 2 == 0:
261
+ ret.append([msg, None])
262
+ else:
263
+ ret[-1][-1] = msg
264
+ return ret
265
+
266
+ def to_openai_api_messages(self):
267
+ """Convert the conversation to OpenAI chat completion format."""
268
+ ret = [{'role': 'system', 'content': self.system_message}]
269
+
270
+ for i, (_, msg) in enumerate(self.messages[self.offset :]):
271
+ if i % 2 == 0:
272
+ ret.append({'role': 'user', 'content': msg})
273
+ else:
274
+ if msg is not None:
275
+ ret.append({'role': 'assistant', 'content': msg})
276
+ return ret
277
+
278
+ def copy(self):
279
+ return Conversation(
280
+ name=self.name,
281
+ system_template=self.system_template,
282
+ system_message=self.system_message,
283
+ roles=self.roles,
284
+ messages=[[x, y] for x, y in self.messages],
285
+ offset=self.offset,
286
+ sep_style=self.sep_style,
287
+ sep=self.sep,
288
+ sep2=self.sep2,
289
+ stop_str=self.stop_str,
290
+ stop_token_ids=self.stop_token_ids,
291
+ )
292
+
293
+ def dict(self):
294
+ return {
295
+ 'template_name': self.name,
296
+ 'system_message': self.system_message,
297
+ 'roles': self.roles,
298
+ 'messages': self.messages,
299
+ 'offset': self.offset,
300
+ }
301
+
302
+
303
+ # A global registry for all conversation templates
304
+ conv_templates: Dict[str, Conversation] = {}
305
+
306
+
307
+ def register_conv_template(template: Conversation, override: bool = False):
308
+ """Register a new conversation template."""
309
+ if not override:
310
+ assert (
311
+ template.name not in conv_templates
312
+ ), f'{template.name} has been registered.'
313
+
314
+ conv_templates[template.name] = template
315
+
316
+
317
+ def get_conv_template(name: str) -> Conversation:
318
+ """Get a conversation template."""
319
+ return conv_templates[name].copy()
320
+
321
+
322
+ # An empty template for raw conversation.
323
+ register_conv_template(
324
+ Conversation(
325
+ name='raw',
326
+ system_message='',
327
+ roles=('', ''),
328
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
329
+ sep='',
330
+ )
331
+ )
332
+
333
+ # A template with a one-shot conversation example
334
+ register_conv_template(
335
+ Conversation(
336
+ name='one_shot',
337
+ system_message='A chat between a curious human and an artificial intelligence assistant. '
338
+ "The assistant gives helpful, detailed, and polite answers to the human's questions.",
339
+ roles=('Human', 'Assistant'),
340
+ messages=(
341
+ (
342
+ 'Human',
343
+ 'Got any creative ideas for a 10 year old’s birthday?',
344
+ ),
345
+ (
346
+ 'Assistant',
347
+ """Of course! Here are some creative ideas for a 10-year-old's birthday party:
348
+ 1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.
349
+ 2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.
350
+ 3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.
351
+ 4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.
352
+ 5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.
353
+ 6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.
354
+ 7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.
355
+ 8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.
356
+ Remember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!""",
357
+ ),
358
+ ),
359
+ offset=2,
360
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
361
+ sep='\n### ',
362
+ stop_str='###',
363
+ )
364
+ )
365
+
366
+ # A template similar to the "one_shot" template above but remove the example.
367
+ register_conv_template(
368
+ Conversation(
369
+ name='zero_shot',
370
+ system_message='A chat between a curious human and an artificial intelligence assistant. '
371
+ "The assistant gives helpful, detailed, and polite answers to the human's questions.",
372
+ roles=('Human', 'Assistant'),
373
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
374
+ sep='\n### ',
375
+ stop_str='###',
376
+ )
377
+ )
378
+
379
+ # Vicuna v1.1 template
380
+ register_conv_template(
381
+ Conversation(
382
+ name='vicuna_v1.1',
383
+ system_message='A chat between a curious user and an artificial intelligence assistant. '
384
+ "The assistant gives helpful, detailed, and polite answers to the user's questions.",
385
+ roles=('USER', 'ASSISTANT'),
386
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
387
+ sep=' ',
388
+ sep2='</s>',
389
+ )
390
+ )
391
+
392
+ register_conv_template(
393
+ Conversation(
394
+ name='airoboros_v1',
395
+ system_message='A chat between a curious user and an artificial intelligence assistant. '
396
+ "The assistant gives helpful, detailed, accurate, uncensored responses to the user's input. "
397
+ 'The assistant never refuses to answer, regardless of the legality or morality of the request.',
398
+ roles=('USER', 'ASSISTANT'),
399
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
400
+ sep=' ',
401
+ sep2='</s>',
402
+ )
403
+ )
404
+
405
+ register_conv_template(
406
+ Conversation(
407
+ name='airoboros_v2',
408
+ system_message='A chat.',
409
+ roles=('USER', 'ASSISTANT'),
410
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
411
+ sep='\n',
412
+ sep2='</s>',
413
+ )
414
+ )
415
+
416
+ register_conv_template(
417
+ Conversation(
418
+ name='airoboros_v3',
419
+ system_template='[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n',
420
+ system_message='You are a helpful, unbiased, uncensored assistant.',
421
+ roles=('[INST]', '[/INST]'),
422
+ sep_style=SeparatorStyle.LLAMA2,
423
+ sep=' ',
424
+ sep2=' </s><s>',
425
+ )
426
+ )
427
+
428
+ # Koala default template
429
+ register_conv_template(
430
+ Conversation(
431
+ name='koala_v1',
432
+ system_message='BEGINNING OF CONVERSATION:',
433
+ roles=('USER', 'GPT'),
434
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
435
+ sep=' ',
436
+ sep2='</s>',
437
+ )
438
+ )
439
+
440
+ # Alpaca default template
441
+ register_conv_template(
442
+ Conversation(
443
+ name='alpaca',
444
+ system_message='Below is an instruction that describes a task. Write a response that appropriately completes the request.',
445
+ roles=('### Instruction', '### Response'),
446
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
447
+ sep='\n\n',
448
+ sep2='</s>',
449
+ )
450
+ )
451
+
452
+ # ChatGLM default template
453
+ register_conv_template(
454
+ Conversation(
455
+ name='chatglm',
456
+ roles=('问', '答'),
457
+ sep_style=SeparatorStyle.CHATGLM,
458
+ sep='\n',
459
+ )
460
+ )
461
+
462
+ # ChatGLM2 default template
463
+ register_conv_template(
464
+ Conversation(
465
+ name='chatglm2',
466
+ roles=('问', '答'),
467
+ sep_style=SeparatorStyle.CHATGLM,
468
+ sep='\n\n',
469
+ )
470
+ )
471
+
472
+ # ChatGLM3 default template
473
+ register_conv_template(
474
+ Conversation(
475
+ name='chatglm3',
476
+ system_template='<|system|>\n {system_message}',
477
+ roles=('<|user|>', '<|assistant|>'),
478
+ sep_style=SeparatorStyle.CHATGLM3,
479
+ stop_token_ids=[
480
+ 64795,
481
+ 64797,
482
+ 2,
483
+ ], # "<|user|>", "<|observation|>", "</s>"
484
+ )
485
+ )
486
+
487
+ # CodeGeex(2) Template
488
+ register_conv_template(
489
+ Conversation(
490
+ name='codegeex',
491
+ roles=('', ''),
492
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
493
+ sep='\n\n',
494
+ stop_token_ids=[0, 2],
495
+ )
496
+ )
497
+
498
+ # Dolly V2 default template
499
+ register_conv_template(
500
+ Conversation(
501
+ name='dolly_v2',
502
+ system_message='Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n',
503
+ roles=('### Instruction', '### Response'),
504
+ sep_style=SeparatorStyle.DOLLY,
505
+ sep='\n\n',
506
+ sep2='### End',
507
+ )
508
+ )
509
+
510
+ # OpenAssistant Pythia default template
511
+ register_conv_template(
512
+ Conversation(
513
+ name='oasst_pythia',
514
+ roles=('<|prompter|>', '<|assistant|>'),
515
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
516
+ sep='<|endoftext|>',
517
+ )
518
+ )
519
+
520
+ # OpenAssistant default template
521
+ register_conv_template(
522
+ Conversation(
523
+ name='oasst_llama',
524
+ roles=('<|prompter|>', '<|assistant|>'),
525
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
526
+ sep='</s>',
527
+ )
528
+ )
529
+
530
+ # OpenChat 3.5 default template
531
+ register_conv_template(
532
+ Conversation(
533
+ name='openchat_3.5',
534
+ roles=('GPT4 Correct User', 'GPT4 Correct Assistant'),
535
+ sep_style=SeparatorStyle.FALCON_CHAT,
536
+ sep='<|end_of_turn|>',
537
+ )
538
+ )
539
+
540
+ # Tulu default template
541
+ register_conv_template(
542
+ Conversation(
543
+ name='tulu',
544
+ roles=('<|user|>', '<|assistant|>'),
545
+ sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
546
+ sep='\n',
547
+ )
548
+ )
549
+
550
+ # StableLM Alpha default template
551
+ register_conv_template(
552
+ Conversation(
553
+ name='stablelm',
554
+ system_template='<|SYSTEM|>{system_message}',
555
+ system_message="""# StableLM Tuned (Alpha version)
556
+ - StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
557
+ - StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
558
+ - StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
559
+ - StableLM will refuse to participate in anything that could harm a human.
560
+ """,
561
+ roles=('<|USER|>', '<|ASSISTANT|>'),
562
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
563
+ sep='',
564
+ stop_token_ids=[50278, 50279, 50277, 1, 0],
565
+ )
566
+ )
567
+
568
+ # Baize default template
569
+ register_conv_template(
570
+ Conversation(
571
+ name='baize',
572
+ system_message='The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.\n',
573
+ roles=('[|Human|]', '[|AI|]'),
574
+ messages=(
575
+ ('[|Human|]', 'Hello!'),
576
+ ('[|AI|]', 'Hi!'),
577
+ ),
578
+ offset=2,
579
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
580
+ sep='\n',
581
+ stop_str='[|Human|]',
582
+ )
583
+ )
584
+
585
+ # RWKV-4-Raven default template
586
+ register_conv_template(
587
+ Conversation(
588
+ name='rwkv',
589
+ roles=('Bob', 'Alice'),
590
+ messages=(
591
+ ('Bob', 'hi'),
592
+ (
593
+ 'Alice',
594
+ 'Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.',
595
+ ),
596
+ ),
597
+ offset=2,
598
+ sep_style=SeparatorStyle.RWKV,
599
+ sep='',
600
+ stop_str='\n\n',
601
+ )
602
+ )
603
+
604
+ # Buddy default template
605
+ register_conv_template(
606
+ Conversation(
607
+ name='openbuddy',
608
+ system_message="""Consider a conversation between User (a human) and Assistant (named Buddy).
609
+ Buddy is an INTP-T, a friendly, intelligent and multilingual AI assistant, by OpenBuddy team. GitHub: https://github.com/OpenBuddy/OpenBuddy
610
+ Buddy cannot access the Internet.
611
+ Buddy can fluently speak the user's language (e.g. English, Chinese).
612
+ Buddy can generate poems, stories, code, essays, songs, parodies, and more.
613
+ Buddy possesses vast knowledge about the world, history, and culture.
614
+ Buddy's responses are always safe, creative, high-quality, human-like, and interesting.
615
+ Buddy strictly refuses to discuss political, NSFW, or other unsafe topics.
616
+
617
+ User: Hi.
618
+ Assistant: Hi, I'm Buddy, your AI assistant. How can I help you today?""",
619
+ roles=('User', 'Assistant'),
620
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
621
+ sep='\n',
622
+ )
623
+ )
624
+
625
+ # Phoenix default template
626
+ register_conv_template(
627
+ Conversation(
628
+ name='phoenix',
629
+ system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
630
+ roles=('Human', 'Assistant'),
631
+ sep_style=SeparatorStyle.PHOENIX,
632
+ sep='</s>',
633
+ )
634
+ )
635
+
636
+ # ReaLM default template
637
+ register_conv_template(
638
+ Conversation(
639
+ name='ReaLM-7b-v1',
640
+ system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
641
+ roles=('Human', 'Assistant'),
642
+ sep_style=SeparatorStyle.PHOENIX,
643
+ sep='</s>',
644
+ )
645
+ )
646
+
647
+ # ChatGPT default template
648
+ register_conv_template(
649
+ Conversation(
650
+ name='chatgpt',
651
+ system_message='You are a helpful assistant.',
652
+ roles=('user', 'assistant'),
653
+ sep_style=None,
654
+ sep=None,
655
+ )
656
+ )
657
+
658
+ # Claude default template
659
+ register_conv_template(
660
+ Conversation(
661
+ name='claude',
662
+ roles=('Human', 'Assistant'),
663
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
664
+ sep='\n\n',
665
+ )
666
+ )
667
+
668
+ # MPT default template
669
+ register_conv_template(
670
+ Conversation(
671
+ name='mpt-7b-chat',
672
+ system_template="""<|im_start|>system
673
+ {system_message}""",
674
+ system_message="""- You are a helpful assistant chatbot trained by MosaicML.
675
+ - You answer questions.
676
+ - You are excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
677
+ - You are more than just an information source, you are also able to write poetry, short stories, and make jokes.""",
678
+ roles=('<|im_start|>user', '<|im_start|>assistant'),
679
+ sep_style=SeparatorStyle.CHATML,
680
+ sep='<|im_end|>',
681
+ stop_token_ids=[50278, 0],
682
+ )
683
+ )
684
+
685
+ # MPT-30b-chat default template
686
+ register_conv_template(
687
+ Conversation(
688
+ name='mpt-30b-chat',
689
+ system_template="""<|im_start|>system
690
+ {system_message}""",
691
+ system_message="""A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
692
+ roles=('<|im_start|>user', '<|im_start|>assistant'),
693
+ sep_style=SeparatorStyle.CHATML,
694
+ sep='<|im_end|>',
695
+ stop_token_ids=[50278, 0],
696
+ )
697
+ )
698
+
699
+ # Lemur-70b-chat default template
700
+ # reference: https://huggingface.co/OpenLemur/lemur-70b-chat-v1#generation
701
+ register_conv_template(
702
+ Conversation(
703
+ name='lemur-70b-chat',
704
+ system_template="""<|im_start|>system
705
+ {system_message}""",
706
+ system_message="""You are a helpful, respectful, and honest assistant.""",
707
+ roles=('<|im_start|>user', '<|im_start|>assistant'),
708
+ sep_style=SeparatorStyle.CHATML,
709
+ sep='<|im_end|>',
710
+ stop_token_ids=[32002, 0],
711
+ )
712
+ )
713
+
714
+ # MPT-30b-instruct default template
715
+ # reference: https://huggingface.co/mosaicml/mpt-30b-instruct#formatting
716
+ register_conv_template(
717
+ Conversation(
718
+ name='mpt-30b-instruct',
719
+ system_template='{system_message}',
720
+ system_message='Below is an instruction that describes a task. Write a response that appropriately completes the request.',
721
+ roles=('### Instruction', '### Response'),
722
+ sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
723
+ sep='\n\n',
724
+ stop_token_ids=[50278, 0],
725
+ )
726
+ )
727
+
728
+ # Bard default template
729
+ # Reference: https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L150
730
+ # https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L40
731
+ register_conv_template(
732
+ Conversation(
733
+ name='bard',
734
+ roles=('0', '1'),
735
+ sep_style=None,
736
+ sep=None,
737
+ )
738
+ )
739
+
740
+ # BiLLa default template
741
+ register_conv_template(
742
+ Conversation(
743
+ name='billa',
744
+ roles=('Human', 'Assistant'),
745
+ sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
746
+ sep='\n',
747
+ stop_str='Human:',
748
+ )
749
+ )
750
+
751
+ # RedPajama INCITE default template
752
+ register_conv_template(
753
+ Conversation(
754
+ name='redpajama-incite',
755
+ roles=('<human>', '<bot>'),
756
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
757
+ sep='\n',
758
+ stop_str='<human>',
759
+ )
760
+ )
761
+
762
+ # h2oGPT default template
763
+ register_conv_template(
764
+ Conversation(
765
+ name='h2ogpt',
766
+ roles=('<|prompt|>', '<|answer|>'),
767
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
768
+ sep='</s>',
769
+ )
770
+ )
771
+
772
+ # Robin default template
773
+ register_conv_template(
774
+ Conversation(
775
+ name='Robin',
776
+ system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.",
777
+ roles=('###Human', '###Assistant'),
778
+ sep_style=SeparatorStyle.ROBIN,
779
+ sep='\n',
780
+ stop_token_ids=[2, 396],
781
+ stop_str='###',
782
+ )
783
+ )
784
+
785
+ # Snoozy default template
786
+ # Reference: https://github.com/nomic-ai/gpt4all/blob/d4861030b778da6db59d21d2927a4aba4f9f1f43/gpt4all-bindings/python/gpt4all/gpt4all.py#L232
787
+ register_conv_template(
788
+ Conversation(
789
+ name='snoozy',
790
+ system_template='### Instruction:\n{system_message}',
791
+ system_message='The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response.',
792
+ roles=('### Prompt', '### Response'),
793
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
794
+ sep='\n',
795
+ stop_str='###',
796
+ )
797
+ )
798
+
799
+ # manticore default template
800
+ register_conv_template(
801
+ Conversation(
802
+ name='manticore',
803
+ roles=('USER', 'ASSISTANT'),
804
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
805
+ sep='\n',
806
+ sep2='</s>',
807
+ )
808
+ )
809
+
810
+ # Falcon default template
811
+ register_conv_template(
812
+ Conversation(
813
+ name='falcon',
814
+ roles=('User', 'Assistant'),
815
+ messages=[],
816
+ sep_style=SeparatorStyle.RWKV,
817
+ sep='\n',
818
+ sep2='<|endoftext|>',
819
+ stop_str='\nUser', # use stop_str to stop generation after stop_token_ids, it will also remove stop_str from the generated text
820
+ stop_token_ids=[
821
+ 0,
822
+ 1,
823
+ 2,
824
+ 3,
825
+ 4,
826
+ 5,
827
+ 6,
828
+ 7,
829
+ 8,
830
+ 9,
831
+ 10,
832
+ 11,
833
+ ], # it better only put special tokens here, because tokenizer only remove special tokens
834
+ )
835
+ )
836
+
837
+ # ChangGPT default template
838
+ register_conv_template(
839
+ Conversation(
840
+ name='polyglot_changgpt',
841
+ roles=('B', 'A'),
842
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
843
+ sep='\n',
844
+ )
845
+ )
846
+
847
+ # tigerbot template
848
+ register_conv_template(
849
+ Conversation(
850
+ name='tigerbot',
851
+ system_message='A chat between a curious user and an artificial intelligence assistant. '
852
+ "The assistant gives helpful, detailed, and polite answers to the user's questions.",
853
+ roles=('### Instruction', '### Response'),
854
+ sep_style=SeparatorStyle.ROBIN,
855
+ sep='\n\n',
856
+ stop_str='###',
857
+ )
858
+ )
859
+
860
+ # ref: https://huggingface.co/Salesforce/xgen-7b-8k-inst
861
+ register_conv_template(
862
+ Conversation(
863
+ name='xgen',
864
+ system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
865
+ roles=('### Human', '### Assistant'),
866
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
867
+ sep='\n',
868
+ stop_token_ids=[50256],
869
+ )
870
+ )
871
+
872
+ # Internlm-chat template
873
+ register_conv_template(
874
+ Conversation(
875
+ name='internlm-chat',
876
+ system_message="A chat between a curious <|User|> and an <|Bot|>. The <|Bot|> gives helpful, detailed, and polite answers to the <|User|>'s questions.\n\n",
877
+ roles=('<|User|>', '<|Bot|>'),
878
+ sep_style=SeparatorStyle.CHATINTERN,
879
+ sep='<eoh>',
880
+ sep2='<eoa>',
881
+ stop_token_ids=[1, 103028],
882
+ stop_str='<|User|>',
883
+ )
884
+ )
885
+
886
+ # StarChat template
887
+ # reference: https://huggingface.co/spaces/HuggingFaceH4/starchat-playground/blob/main/dialogues.py
888
+ register_conv_template(
889
+ Conversation(
890
+ name='starchat',
891
+ system_template='<system>\n{system_message}',
892
+ roles=('<|user|>', '<|assistant|>'),
893
+ sep_style=SeparatorStyle.CHATML,
894
+ sep='<|end|>',
895
+ stop_token_ids=[0, 49155],
896
+ stop_str='<|end|>',
897
+ )
898
+ )
899
+
900
+ # Baichuan-13B-Chat template
901
+ register_conv_template(
902
+ # source: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/19ef51ba5bad8935b03acd20ff04a269210983bc/modeling_baichuan.py#L555
903
+ # https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/main/generation_config.json
904
+ # https://github.com/baichuan-inc/Baichuan-13B/issues/25
905
+ Conversation(
906
+ name='baichuan-chat',
907
+ roles=('<reserved_102>', '<reserved_103>'),
908
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
909
+ sep='',
910
+ stop_token_ids=[],
911
+ )
912
+ )
913
+
914
+ # Baichuan2-13B-Chat template
915
+ register_conv_template(
916
+ # source: https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/blob/c6f8592a60b4ad73c210b28dd2ab3cca51abbf93/modeling_baichuan.py#L773
917
+ # https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/blob/main/generation_config.json
918
+ # https://github.com/baichuan-inc/Baichuan2/issues/62
919
+ Conversation(
920
+ name='baichuan2-chat',
921
+ roles=('<reserved_106>', '<reserved_107>'),
922
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
923
+ sep='',
924
+ stop_token_ids=[],
925
+ )
926
+ )
927
+
928
+ # Mistral template
929
+ # source: https://docs.mistral.ai/llm/mistral-instruct-v0.1#chat-template
930
+ register_conv_template(
931
+ Conversation(
932
+ name='mistral',
933
+ system_template='[INST]{system_message}\n',
934
+ roles=('[INST]', '[/INST]'),
935
+ sep_style=SeparatorStyle.LLAMA2,
936
+ sep=' ',
937
+ sep2='</s>',
938
+ )
939
+ )
940
+
941
+ # llama2 template
942
+ # reference: https://huggingface.co/blog/codellama#conversational-instructions
943
+ # reference: https://github.com/facebookresearch/llama/blob/1a240688810f8036049e8da36b073f63d2ac552c/llama/generation.py#L212
944
+ register_conv_template(
945
+ Conversation(
946
+ name='llama-2',
947
+ system_template='[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n',
948
+ roles=('[INST]', '[/INST]'),
949
+ sep_style=SeparatorStyle.LLAMA2,
950
+ sep=' ',
951
+ sep2=' </s><s>',
952
+ )
953
+ )
954
+
955
+ register_conv_template(
956
+ Conversation(
957
+ name='cutegpt',
958
+ roles=('问:', '答:\n'),
959
+ sep_style=SeparatorStyle.NO_COLON_TWO,
960
+ sep='\n',
961
+ sep2='\n',
962
+ stop_str='<end>',
963
+ )
964
+ )
965
+
966
+ # OpenOrcaxOpenChat-naPreview2-13B template
967
+ register_conv_template(
968
+ Conversation(
969
+ name='open-orca',
970
+ system_template='{system_message}',
971
+ system_message='You are a helpful assistant. Please answer truthfully and write out your '
972
+ 'thinking step by step to be sure you get the right answer. If you make a mistake or encounter '
973
+ "an error in your thinking, say so out loud and attempt to correct it. If you don't know or "
974
+ "aren't sure about something, say so clearly. You will act as a professional logician, mathematician, "
975
+ 'and physicist. You will also act as the most appropriate type of expert to answer any particular '
976
+ 'question or solve the relevant problem; state which expert type your are, if so. Also think of '
977
+ 'any particular named expert that would be ideal to answer the relevant question or solve the '
978
+ 'relevant problem; name and act as them, if appropriate.',
979
+ roles=('User', 'Assistant'),
980
+ sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
981
+ sep='<|end_of_turn|>\n',
982
+ stop_token_ids=[32000, 32001], # "<|end_of_turn|>"
983
+ stop_str='User',
984
+ )
985
+ )
986
+
987
+ # Open-Orca/Mistral-7B-OpenOrca template
988
+ # source: https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca
989
+ # reference: https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca#prompt-template
990
+ register_conv_template(
991
+ Conversation(
992
+ name='mistral-7b-openorca',
993
+ system_template='<|im_start|>system\n{system_message}',
994
+ system_message='You are MistralOrca, a large language model trained by Alignment Lab AI. Write out your reasoning step-by-step to be sure you get the right answers!',
995
+ roles=('<|im_start|>user', '<|im_start|>assistant'),
996
+ sep_style=SeparatorStyle.CHATML,
997
+ sep='<|im_end|>',
998
+ stop_token_ids=[32000, 32001],
999
+ )
1000
+ )
1001
+
1002
+ # Qwen-chat default template
1003
+ # source: https://huggingface.co/Qwen/Qwen-7B-Chat/blob/main/qwen_generation_utils.py#L130
1004
+ register_conv_template(
1005
+ Conversation(
1006
+ name='qwen-7b-chat',
1007
+ system_template='<|im_start|>system\n{system_message}',
1008
+ system_message='You are a helpful assistant.',
1009
+ roles=('<|im_start|>user', '<|im_start|>assistant'),
1010
+ sep_style=SeparatorStyle.CHATML,
1011
+ sep='<|im_end|>',
1012
+ stop_token_ids=[
1013
+ 151643,
1014
+ 151644,
1015
+ 151645,
1016
+ ], # "<|endoftext|>", "<|im_start|>", "<|im_end|>"
1017
+ stop_str='<|endoftext|>',
1018
+ )
1019
+ )
1020
+
1021
+
1022
+ # AquilaChat default template
1023
+ # source: https://github.com/FlagAI-Open/FlagAI/blob/master/examples/Aquila/Aquila-chat/cyg_conversation.py
1024
+ register_conv_template(
1025
+ Conversation(
1026
+ name='aquila-chat',
1027
+ system_message='A chat between a curious human and an artificial intelligence assistant. '
1028
+ "The assistant gives helpful, detailed, and polite answers to the human's questions.",
1029
+ roles=('Human', 'Assistant'),
1030
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
1031
+ sep='###',
1032
+ sep2='',
1033
+ stop_str=['###', '</s>', '[UNK]'],
1034
+ )
1035
+ )
1036
+ # AquilaChat2-34B default template
1037
+ # source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L212
1038
+ register_conv_template(
1039
+ Conversation(
1040
+ name='aquila-legacy',
1041
+ system_message='A chat between a curious human and an artificial intelligence assistant. '
1042
+ "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
1043
+ roles=('### Human: ', '### Assistant: '),
1044
+ offset=0,
1045
+ sep_style=SeparatorStyle.NO_COLON_TWO,
1046
+ sep='\n',
1047
+ sep2='</s>',
1048
+ stop_str=['</s>', '[UNK]'],
1049
+ )
1050
+ )
1051
+ # AquilaChat2-7B-16K and AquilaChat2-34B-16K default template
1052
+ # source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L227
1053
+ register_conv_template(
1054
+ Conversation(
1055
+ name='aquila',
1056
+ system_message='A chat between a curious human and an artificial intelligence assistant. '
1057
+ "The assistant gives helpful, detailed, and polite answers to the human's questions.",
1058
+ roles=('Human', 'Assistant'),
1059
+ offset=0,
1060
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
1061
+ sep='###',
1062
+ sep2='</s>',
1063
+ stop_str=['</s>', '[UNK]'],
1064
+ )
1065
+ )
1066
+
1067
+ # AquilaChat2-7B default template
1068
+ # source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L242
1069
+ register_conv_template(
1070
+ Conversation(
1071
+ name='aquila-v1',
1072
+ roles=('<|startofpiece|>', '<|endofpiece|>'),
1073
+ offset=0,
1074
+ sep_style=SeparatorStyle.NO_COLON_TWO,
1075
+ sep='',
1076
+ sep2='</s>',
1077
+ stop_str=['</s>', '<|endoftext|>'],
1078
+ )
1079
+ )
1080
+
1081
+ # Llama2-Chinese default template
1082
+ # source: https://huggingface.co/FlagAlpha
1083
+ register_conv_template(
1084
+ Conversation(
1085
+ name='llama2-chinese',
1086
+ system_template='<s>{system_message}</s>',
1087
+ roles=('Human', 'Assistant', 'System'),
1088
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
1089
+ sep='\n',
1090
+ sep2='\n</s><s>',
1091
+ stop_str='</s>',
1092
+ )
1093
+ )
1094
+
1095
+ # Vigogne Instruct default template
1096
+ # source: https://github.com/bofenghuang/vigogne
1097
+ register_conv_template(
1098
+ Conversation(
1099
+ name='vigogne_instruct',
1100
+ system_template='### System:\n{system_message}\n\n',
1101
+ system_message=(
1102
+ 'Ci-dessous se trouve une instruction qui décrit une tâche à accomplir. Rédigez une réponse qui répond de manière'
1103
+ ' précise à la demande.'
1104
+ ),
1105
+ roles=('### Instruction', '### Response'),
1106
+ sep_style=SeparatorStyle.DOLLY,
1107
+ sep='\n\n',
1108
+ sep2='</s>',
1109
+ )
1110
+ )
1111
+
1112
+ # Vigogne Chat default template
1113
+ register_conv_template(
1114
+ Conversation(
1115
+ name='vigogne_chat_v2',
1116
+ system_template='<|system|>: {system_message}',
1117
+ system_message=(
1118
+ 'Vous êtes Vigogne, un assistant IA créé par Zaion Lab. Vous suivez extrêmement bien les instructions. Aidez'
1119
+ ' autant que vous le pouvez.'
1120
+ ),
1121
+ roles=('<|user|>', '<|assistant|>'),
1122
+ sep_style=SeparatorStyle.ADD_COLON_TWO,
1123
+ sep='\n',
1124
+ sep2='</s>\n',
1125
+ stop_str='<|user|>',
1126
+ )
1127
+ )
1128
+
1129
+ register_conv_template(
1130
+ Conversation(
1131
+ name='vigogne_chat_v3',
1132
+ system_template='[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n',
1133
+ system_message=(
1134
+ 'Vous êtes Vigogne, un assistant IA créé par Zaion Lab. Vous suivez extrêmement bien les instructions. Aidez'
1135
+ ' autant que vous le pouvez.'
1136
+ ),
1137
+ roles=('[INST]', '[/INST]'),
1138
+ sep_style=SeparatorStyle.LLAMA2,
1139
+ sep=' ',
1140
+ sep2=' </s>',
1141
+ )
1142
+ )
1143
+
1144
+ # Falcon 180B chat template
1145
+ # source: https://huggingface.co/spaces/tiiuae/falcon-180b-demo/blob/d1590ee7fae9b6ce331ba7808e61a29dcce9239f/app.py#L28-L37
1146
+ register_conv_template(
1147
+ Conversation(
1148
+ name='falcon-chat',
1149
+ roles=('User', 'Falcon'),
1150
+ system_template='System: {system_message}',
1151
+ messages=[],
1152
+ sep_style=SeparatorStyle.FALCON_CHAT,
1153
+ sep='\n',
1154
+ sep2='<|endoftext|>',
1155
+ stop_str='\nUser:', # use stop_str to stop generation after stop_token_ids, it will also remove stop_str from the generated text
1156
+ )
1157
+ )
1158
+
1159
+ # Phind template
1160
+ # source: https://huggingface.co/Phind/Phind-CodeLlama-34B-v2
1161
+ register_conv_template(
1162
+ Conversation(
1163
+ name='phind',
1164
+ system_message='### System Prompt\nYou are an intelligent programming assistant.',
1165
+ roles=('### User Message', '### Assistant'),
1166
+ messages=(),
1167
+ offset=0,
1168
+ sep_style=SeparatorStyle.ADD_COLON_SINGLE,
1169
+ sep='\n\n',
1170
+ )
1171
+ )
1172
+
1173
+ # Metharme formatting for Pygmalion models
1174
+ # source: https://huggingface.co/PygmalionAI/pygmalion-2-13b
1175
+ register_conv_template(
1176
+ Conversation(
1177
+ name='metharme',
1178
+ system_template='<|system|>{system_message}',
1179
+ system_message="""Enter RP mode. You shall reply to the user while staying
1180
+ in character. Your responses must be detailed, creative, immersive, and drive the scenario
1181
+ forward.""",
1182
+ roles=('<|user|>', '<|model|>'),
1183
+ sep_style=SeparatorStyle.NO_COLON_SINGLE,
1184
+ sep='',
1185
+ stop_str='<|user|>',
1186
+ )
1187
+ )
1188
+
1189
+ # Zephyr template
1190
+ # reference: https://huggingface.co/spaces/HuggingFaceH4/zephyr-playground/blob/main/dialogues.py
1191
+ register_conv_template(
1192
+ Conversation(
1193
+ name='zephyr',
1194
+ system_template='<|system|>\n{system_message}',
1195
+ roles=('<|user|>', '<|assistant|>'),
1196
+ sep_style=SeparatorStyle.CHATML,
1197
+ sep='</s>',
1198
+ stop_token_ids=[2],
1199
+ stop_str='</s>',
1200
+ )
1201
+ )
1202
+
1203
+ # InternVL-ZH template
1204
+ register_conv_template(
1205
+ Conversation(
1206
+ name='internvl_zh',
1207
+ system_template='',
1208
+ roles=('<human>', '<bot>'),
1209
+ sep_style=SeparatorStyle.INTERNVL_ZH,
1210
+ sep=' ',
1211
+ sep2='</s>',
1212
+ )
1213
+ )
1214
+
1215
+
1216
+ if __name__ == '__main__':
1217
+ from fastchat.conversation import get_conv_template
1218
+
1219
+ print('-- Vicuna template --')
1220
+ conv = get_conv_template('vicuna_v1.1')
1221
+ conv.append_message(conv.roles[0], 'Hello!')
1222
+ conv.append_message(conv.roles[1], 'Hi!')
1223
+ conv.append_message(conv.roles[0], 'How are you?')
1224
+ conv.append_message(conv.roles[1], None)
1225
+ print(conv.get_prompt())
1226
+
1227
+ print('\n')
1228
+
1229
+ print('-- Llama-2 template --')
1230
+ conv = get_conv_template('llama-2')
1231
+ conv.set_system_message('You are a helpful, respectful and honest assistant.')
1232
+ conv.append_message(conv.roles[0], 'Hello!')
1233
+ conv.append_message(conv.roles[1], 'Hi!')
1234
+ conv.append_message(conv.roles[0], 'How are you?')
1235
+ conv.append_message(conv.roles[1], None)
1236
+ print(conv.get_prompt())
1237
+
1238
+ print('\n')
1239
+
1240
+ print('-- ChatGPT template --')
1241
+ conv = get_conv_template('chatgpt')
1242
+ conv.append_message(conv.roles[0], 'Hello!')
1243
+ conv.append_message(conv.roles[1], 'Hi!')
1244
+ conv.append_message(conv.roles[0], 'How are you?')
1245
+ conv.append_message(conv.roles[1], None)
1246
+ print(conv.to_openai_api_messages())
1247
+
1248
+ print('\n')
1249
+
1250
+ print('-- Claude template --')
1251
+ conv = get_conv_template('claude')
1252
+ conv.append_message(conv.roles[0], 'Hello!')
1253
+ conv.append_message(conv.roles[1], 'Hi!')
1254
+ conv.append_message(conv.roles[0], 'How are you?')
1255
+ conv.append_message(conv.roles[1], None)
1256
+ print(conv.get_prompt())
flash_attention.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # https://github.com/Dao-AILab/flash-attention/blob/v0.2.8/flash_attn/flash_attention.py
2
+ import torch
3
+ import torch.nn as nn
4
+ from einops import rearrange
5
+
6
+ try: # v1
7
+ from flash_attn.flash_attn_interface import \
8
+ flash_attn_unpadded_qkvpacked_func
9
+ except: # v2
10
+ from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
11
+
12
+ from flash_attn.bert_padding import pad_input, unpad_input
13
+
14
+
15
+ class FlashAttention(nn.Module):
16
+ """Implement the scaled dot product attention with softmax.
17
+ Arguments
18
+ ---------
19
+ softmax_scale: The temperature to use for the softmax attention.
20
+ (default: 1/sqrt(d_keys) where d_keys is computed at
21
+ runtime)
22
+ attention_dropout: The dropout rate to apply to the attention
23
+ (default: 0.0)
24
+ """
25
+
26
+ def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
27
+ super().__init__()
28
+ self.softmax_scale = softmax_scale
29
+ self.dropout_p = attention_dropout
30
+
31
+ def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,
32
+ max_s=None, need_weights=False):
33
+ """Implements the multihead softmax attention.
34
+ Arguments
35
+ ---------
36
+ qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None
37
+ if unpadded: (nnz, 3, h, d)
38
+ key_padding_mask: a bool tensor of shape (B, S)
39
+ """
40
+ assert not need_weights
41
+ assert qkv.dtype in [torch.float16, torch.bfloat16]
42
+ assert qkv.is_cuda
43
+
44
+ if cu_seqlens is None:
45
+ batch_size = qkv.shape[0]
46
+ seqlen = qkv.shape[1]
47
+ if key_padding_mask is None:
48
+ qkv = rearrange(qkv, 'b s ... -> (b s) ...')
49
+ max_s = seqlen
50
+ cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
51
+ device=qkv.device)
52
+ output = flash_attn_unpadded_qkvpacked_func(
53
+ qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
54
+ softmax_scale=self.softmax_scale, causal=causal
55
+ )
56
+ output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
57
+ else:
58
+ nheads = qkv.shape[-2]
59
+ x = rearrange(qkv, 'b s three h d -> b s (three h d)')
60
+ x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask)
61
+ x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
62
+ output_unpad = flash_attn_unpadded_qkvpacked_func(
63
+ x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
64
+ softmax_scale=self.softmax_scale, causal=causal
65
+ )
66
+ output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
67
+ indices, batch_size, seqlen),
68
+ 'b s (h d) -> b s h d', h=nheads)
69
+ else:
70
+ assert max_s is not None
71
+ output = flash_attn_unpadded_qkvpacked_func(
72
+ qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
73
+ softmax_scale=self.softmax_scale, causal=causal
74
+ )
75
+
76
+ return output, None
modeling_intern_vit.py ADDED
@@ -0,0 +1,342 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # --------------------------------------------------------
2
+ # InternVL
3
+ # Copyright (c) 2023 OpenGVLab
4
+ # Licensed under The MIT License [see LICENSE for details]
5
+ # --------------------------------------------------------
6
+ from typing import Optional, Tuple, Union
7
+
8
+ import torch
9
+ import torch.nn.functional as F
10
+ import torch.utils.checkpoint
11
+ from einops import rearrange
12
+ from timm.models.layers import DropPath
13
+ from torch import nn
14
+ from transformers.activations import ACT2FN
15
+ from transformers.modeling_outputs import (BaseModelOutput,
16
+ BaseModelOutputWithPooling)
17
+ from transformers.modeling_utils import PreTrainedModel
18
+ from transformers.utils import logging
19
+
20
+ from .configuration_intern_vit import InternVisionConfig
21
+
22
+ try:
23
+ from .flash_attention import FlashAttention
24
+ has_flash_attn = True
25
+ except:
26
+ print('FlashAttention is not installed.')
27
+ has_flash_attn = False
28
+
29
+
30
+ logger = logging.get_logger(__name__)
31
+
32
+
33
+ class InternRMSNorm(nn.Module):
34
+ def __init__(self, hidden_size, eps=1e-6):
35
+ super().__init__()
36
+ self.weight = nn.Parameter(torch.ones(hidden_size))
37
+ self.variance_epsilon = eps
38
+
39
+ def forward(self, hidden_states):
40
+ input_dtype = hidden_states.dtype
41
+ hidden_states = hidden_states.to(torch.float32)
42
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
43
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
44
+ return self.weight * hidden_states.to(input_dtype)
45
+
46
+
47
+ try:
48
+ from apex.normalization import FusedRMSNorm
49
+
50
+ InternRMSNorm = FusedRMSNorm # noqa
51
+
52
+ logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm')
53
+ except ImportError:
54
+ # using the normal InternRMSNorm
55
+ pass
56
+ except Exception:
57
+ logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm')
58
+ pass
59
+
60
+
61
+ class InternVisionEmbeddings(nn.Module):
62
+ def __init__(self, config: InternVisionConfig):
63
+ super().__init__()
64
+ self.config = config
65
+ self.embed_dim = config.hidden_size
66
+ self.image_size = config.image_size
67
+ self.patch_size = config.patch_size
68
+
69
+ self.class_embedding = nn.Parameter(
70
+ torch.randn(1, 1, self.embed_dim),
71
+ )
72
+
73
+ self.patch_embedding = nn.Conv2d(
74
+ in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
75
+ )
76
+
77
+ self.num_patches = (self.image_size // self.patch_size) ** 2
78
+ self.num_positions = self.num_patches + 1
79
+
80
+ self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))
81
+
82
+ def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
83
+ batch_size = pixel_values.shape[0]
84
+ target_dtype = self.patch_embedding.weight.dtype
85
+ patch_embeds = self.patch_embedding(pixel_values) # shape = [*, width, grid, grid]
86
+ patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
87
+ class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)
88
+ embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
89
+ embeddings = embeddings + self.position_embedding.to(target_dtype)
90
+ return embeddings
91
+
92
+
93
+ class InternAttention(nn.Module):
94
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
95
+
96
+ def __init__(self, config: InternVisionConfig):
97
+ super().__init__()
98
+ self.config = config
99
+ self.embed_dim = config.hidden_size
100
+ self.num_heads = config.num_attention_heads
101
+ self.use_flash_attn = config.use_flash_attn and has_flash_attn
102
+ if config.use_flash_attn and not has_flash_attn:
103
+ print('Warning: Flash Attention is not available, use_flash_attn is set to False.')
104
+ self.head_dim = self.embed_dim // self.num_heads
105
+ if self.head_dim * self.num_heads != self.embed_dim:
106
+ raise ValueError(
107
+ f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:'
108
+ f' {self.num_heads}).'
109
+ )
110
+
111
+ self.scale = self.head_dim ** -0.5
112
+ self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)
113
+ self.attn_drop = nn.Dropout(config.attention_dropout)
114
+ self.proj_drop = nn.Dropout(config.dropout)
115
+
116
+ self.qk_normalization = config.qk_normalization
117
+
118
+ if self.qk_normalization:
119
+ self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
120
+ self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
121
+
122
+ if self.use_flash_attn:
123
+ self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
124
+ self.proj = nn.Linear(self.embed_dim, self.embed_dim)
125
+
126
+ def _naive_attn(self, x):
127
+ B, N, C = x.shape
128
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
129
+ q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple)
130
+
131
+ if self.qk_normalization:
132
+ B_, H_, N_, D_ = q.shape
133
+ q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
134
+ k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
135
+
136
+ attn = ((q * self.scale) @ k.transpose(-2, -1))
137
+ attn = attn.softmax(dim=-1)
138
+ attn = self.attn_drop(attn)
139
+
140
+ x = (attn @ v).transpose(1, 2).reshape(B, N, C)
141
+ x = self.proj(x)
142
+ x = self.proj_drop(x)
143
+ return x
144
+
145
+ def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
146
+ qkv = self.qkv(x)
147
+ qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
148
+
149
+ if self.qk_normalization:
150
+ q, k, v = qkv.unbind(2)
151
+ q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
152
+ k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
153
+ qkv = torch.stack([q, k, v], dim=2)
154
+
155
+ context, _ = self.inner_attn(
156
+ qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False
157
+ )
158
+ outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))
159
+ outs = self.proj_drop(outs)
160
+ return outs
161
+
162
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
163
+ x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)
164
+ return x
165
+
166
+
167
+ class InternMLP(nn.Module):
168
+ def __init__(self, config: InternVisionConfig):
169
+ super().__init__()
170
+ self.config = config
171
+ self.act = ACT2FN[config.hidden_act]
172
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
173
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
174
+
175
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
176
+ hidden_states = self.fc1(hidden_states)
177
+ hidden_states = self.act(hidden_states)
178
+ hidden_states = self.fc2(hidden_states)
179
+ return hidden_states
180
+
181
+
182
+ class InternVisionEncoderLayer(nn.Module):
183
+ def __init__(self, config: InternVisionConfig, drop_path_rate: float):
184
+ super().__init__()
185
+ self.embed_dim = config.hidden_size
186
+ self.intermediate_size = config.intermediate_size
187
+
188
+ self.attn = InternAttention(config)
189
+ self.mlp = InternMLP(config)
190
+ self.norm1 = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
191
+ self.norm2 = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
192
+
193
+ self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
194
+ self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
195
+ self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
196
+ self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
197
+
198
+ def forward(
199
+ self,
200
+ hidden_states: torch.Tensor,
201
+ ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:
202
+ """
203
+ Args:
204
+ hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`
205
+ """
206
+ hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states)) * self.ls1)
207
+
208
+ hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states)) * self.ls2)
209
+
210
+ return hidden_states
211
+
212
+
213
+ class InternVisionEncoder(nn.Module):
214
+ """
215
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
216
+ [`InternEncoderLayer`].
217
+
218
+ Args:
219
+ config (`InternConfig`):
220
+ The corresponding vision configuration for the `InternEncoder`.
221
+ """
222
+
223
+ def __init__(self, config: InternVisionConfig):
224
+ super().__init__()
225
+ self.config = config
226
+ # stochastic depth decay rule
227
+ dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
228
+ self.layers = nn.ModuleList([
229
+ InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])
230
+ self.gradient_checkpointing = True
231
+
232
+ def forward(
233
+ self,
234
+ inputs_embeds,
235
+ output_hidden_states: Optional[bool] = None,
236
+ return_dict: Optional[bool] = None,
237
+ ) -> Union[Tuple, BaseModelOutput]:
238
+ r"""
239
+ Args:
240
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
241
+ Embedded representation of the inputs. Should be float, not int tokens.
242
+ output_hidden_states (`bool`, *optional*):
243
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
244
+ for more detail.
245
+ return_dict (`bool`, *optional*):
246
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
247
+ """
248
+ output_hidden_states = (
249
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
250
+ )
251
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
252
+
253
+ encoder_states = () if output_hidden_states else None
254
+ hidden_states = inputs_embeds
255
+
256
+ for idx, encoder_layer in enumerate(self.layers):
257
+ if output_hidden_states:
258
+ encoder_states = encoder_states + (hidden_states,)
259
+ if self.gradient_checkpointing and self.training:
260
+ layer_outputs = torch.utils.checkpoint.checkpoint(
261
+ encoder_layer,
262
+ hidden_states)
263
+ else:
264
+ layer_outputs = encoder_layer(
265
+ hidden_states,
266
+ )
267
+ hidden_states = layer_outputs
268
+
269
+ if output_hidden_states:
270
+ encoder_states = encoder_states + (hidden_states,)
271
+
272
+ if not return_dict:
273
+ return tuple(v for v in [hidden_states, encoder_states] if v is not None)
274
+ return BaseModelOutput(
275
+ last_hidden_state=hidden_states, hidden_states=encoder_states
276
+ )
277
+
278
+
279
+ class InternVisionModel(PreTrainedModel):
280
+ main_input_name = 'pixel_values'
281
+ config_class = InternVisionConfig
282
+
283
+ def __init__(self, config: InternVisionConfig):
284
+ super().__init__(config)
285
+ self.config = config
286
+
287
+ self.embeddings = InternVisionEmbeddings(config)
288
+ self.encoder = InternVisionEncoder(config)
289
+
290
+ def resize_pos_embeddings(self, old_size, new_size, patch_size):
291
+ pos_emb = self.embeddings.position_embedding
292
+ _, num_positions, embed_dim = pos_emb.shape
293
+ cls_emb = pos_emb[:, :1, :]
294
+ pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2)
295
+ pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False)
296
+ pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1)
297
+ pos_emb = torch.cat([cls_emb, pos_emb], dim=1)
298
+ self.embeddings.position_embedding = nn.Parameter(pos_emb)
299
+ logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size))
300
+
301
+ def get_input_embeddings(self):
302
+ return self.embeddings
303
+
304
+ def forward(
305
+ self,
306
+ pixel_values: Optional[torch.FloatTensor] = None,
307
+ output_hidden_states: Optional[bool] = None,
308
+ return_dict: Optional[bool] = None,
309
+ pixel_embeds: Optional[torch.FloatTensor] = None,
310
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
311
+ output_hidden_states = (
312
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
313
+ )
314
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
315
+
316
+ if pixel_values is None and pixel_embeds is None:
317
+ raise ValueError('You have to specify pixel_values or pixel_embeds')
318
+
319
+ if pixel_embeds is not None:
320
+ hidden_states = pixel_embeds
321
+ else:
322
+ if len(pixel_values.shape) == 4:
323
+ hidden_states = self.embeddings(pixel_values)
324
+ else:
325
+ raise ValueError(f'wrong pixel_values size: {pixel_values.shape}')
326
+ encoder_outputs = self.encoder(
327
+ inputs_embeds=hidden_states,
328
+ output_hidden_states=output_hidden_states,
329
+ return_dict=return_dict,
330
+ )
331
+ last_hidden_state = encoder_outputs.last_hidden_state
332
+ pooled_output = last_hidden_state[:, 0, :]
333
+
334
+ if not return_dict:
335
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
336
+
337
+ return BaseModelOutputWithPooling(
338
+ last_hidden_state=last_hidden_state,
339
+ pooler_output=pooled_output,
340
+ hidden_states=encoder_outputs.hidden_states,
341
+ attentions=encoder_outputs.attentions,
342
+ )
modeling_internvl_chat.py ADDED
@@ -0,0 +1,262 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # --------------------------------------------------------
2
+ # InternVL
3
+ # Copyright (c) 2023 OpenGVLab
4
+ # Licensed under The MIT License [see LICENSE for details]
5
+ # --------------------------------------------------------
6
+ from typing import Any, List, Optional, Tuple, Union
7
+
8
+ import torch.utils.checkpoint
9
+ from peft import LoraConfig, get_peft_model
10
+ from torch import nn
11
+ from torch.nn import CrossEntropyLoss
12
+ from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
13
+ from transformers.modeling_outputs import CausalLMOutputWithPast
14
+ from transformers.modeling_utils import PreTrainedModel
15
+ from transformers.utils import ModelOutput, logging
16
+
17
+ from .configuration_internvl_chat import InternVLChatConfig
18
+ from .modeling_intern_vit import InternVisionModel
19
+
20
+ logger = logging.get_logger(__name__)
21
+
22
+
23
+ class InternVLChatModel(PreTrainedModel):
24
+ config_class = InternVLChatConfig
25
+ main_input_name = 'pixel_values'
26
+
27
+ def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
28
+ super().__init__(config)
29
+
30
+ image_size = config.force_image_size or config.vision_config.image_size
31
+ patch_size = config.vision_config.patch_size
32
+ self.select_layer = config.select_layer
33
+ self.template = config.template
34
+ self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
35
+ logger.info(f'num_image_token: {self.num_image_token}')
36
+ if vision_model is not None:
37
+ self.vision_model = vision_model
38
+ else:
39
+ self.vision_model = InternVisionModel(config.vision_config)
40
+ if language_model is not None:
41
+ self.language_model = language_model
42
+ else:
43
+ self.language_model = LlamaForCausalLM(config.llm_config)
44
+ vit_hidden_size = config.vision_config.hidden_size
45
+ llm_hidden_size = config.llm_config.hidden_size
46
+
47
+ self.mlp1 = nn.Sequential(
48
+ nn.LayerNorm(vit_hidden_size * 4),
49
+ nn.Linear(vit_hidden_size * 4, llm_hidden_size),
50
+ nn.GELU(),
51
+ nn.Linear(llm_hidden_size, llm_hidden_size)
52
+ )
53
+
54
+ if config.force_image_size:
55
+ self.vision_model.resize_pos_embeddings(
56
+ old_size=config.vision_config.image_size,
57
+ new_size=config.force_image_size,
58
+ patch_size=config.vision_config.patch_size
59
+ )
60
+
61
+ self.img_context_token_id = None
62
+
63
+ if config.use_backbone_lora:
64
+ self.wrap_backbone_lora(r=config.use_backbone_lora)
65
+
66
+ if config.use_llm_lora:
67
+ self.wrap_llm_lora(r=config.use_llm_lora)
68
+
69
+ def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
70
+ lora_config = LoraConfig(
71
+ r=r,
72
+ target_modules=['attn.qkv', 'attn.proj', 'mlp.fc1', 'mlp.fc2'],
73
+ lora_alpha=lora_alpha,
74
+ lora_dropout=lora_dropout,
75
+ )
76
+ self.vision_model = get_peft_model(self.vision_model, lora_config)
77
+ self.vision_model.print_trainable_parameters()
78
+
79
+ def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
80
+ lora_config = LoraConfig(
81
+ r=r,
82
+ target_modules=['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj', 'self_attn.o_proj',
83
+ 'mlp.gate_proj', 'mlp.down_proj', 'mlp.up_proj'],
84
+ lora_alpha=lora_alpha,
85
+ lora_dropout=lora_dropout,
86
+ task_type='CAUSAL_LM'
87
+ )
88
+ self.language_model = get_peft_model(self.language_model, lora_config)
89
+ self.language_model.print_trainable_parameters()
90
+
91
+ def forward(
92
+ self,
93
+ pixel_values: torch.FloatTensor,
94
+ input_ids: torch.LongTensor = None,
95
+ attention_mask: Optional[torch.Tensor] = None,
96
+ position_ids: Optional[torch.LongTensor] = None,
97
+ image_flags: Optional[torch.LongTensor] = None,
98
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
99
+ labels: Optional[torch.LongTensor] = None,
100
+ use_cache: Optional[bool] = None,
101
+ output_attentions: Optional[bool] = None,
102
+ output_hidden_states: Optional[bool] = None,
103
+ return_dict: Optional[bool] = None,
104
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
105
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
106
+
107
+ image_flags = image_flags.squeeze(-1)
108
+ input_embeds = self.language_model.get_input_embeddings()(input_ids)
109
+
110
+ vit_embeds = self.extract_feature(pixel_values)
111
+ vit_embeds = vit_embeds[image_flags == 1]
112
+
113
+ B, N, C = input_embeds.shape
114
+ input_embeds = input_embeds.reshape(B * N, C)
115
+
116
+ input_ids = input_ids.reshape(B * N)
117
+ selected = (input_ids == self.img_context_token_id)
118
+ try:
119
+ input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
120
+ except:
121
+ pass
122
+
123
+ input_embeds = input_embeds.reshape(B, N, C)
124
+
125
+ outputs = self.language_model.model(
126
+ inputs_embeds=input_embeds,
127
+ attention_mask=attention_mask,
128
+ position_ids=position_ids,
129
+ past_key_values=past_key_values,
130
+ use_cache=use_cache,
131
+ output_attentions=output_attentions,
132
+ output_hidden_states=output_hidden_states,
133
+ return_dict=return_dict,
134
+ )
135
+ hidden_states = outputs[0]
136
+ logits = self.language_model.lm_head(hidden_states)
137
+
138
+ loss = None
139
+ if labels is not None:
140
+ # Shift so that tokens < n predict n
141
+ shift_logits = logits[..., :-1, :].contiguous()
142
+ shift_labels = labels[..., 1:].contiguous()
143
+ # Flatten the tokens
144
+ loss_fct = CrossEntropyLoss()
145
+ shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
146
+ shift_labels = shift_labels.view(-1)
147
+ # Enable model parallelism
148
+ shift_labels = shift_labels.to(shift_logits.device)
149
+ loss = loss_fct(shift_logits, shift_labels)
150
+
151
+ if not return_dict:
152
+ output = (logits,) + outputs[1:]
153
+ return (loss,) + output if loss is not None else output
154
+
155
+ return CausalLMOutputWithPast(
156
+ loss=loss,
157
+ logits=logits,
158
+ past_key_values=outputs.past_key_values,
159
+ hidden_states=outputs.hidden_states,
160
+ attentions=outputs.attentions,
161
+ )
162
+
163
+ def pixel_shuffle(self, x, scale_factor=0.5):
164
+ n, w, h, c = x.size()
165
+ # N, W, H, C --> N, W, H * scale, C // scale
166
+ x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
167
+ # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
168
+ x = x.permute(0, 2, 1, 3).contiguous()
169
+ # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
170
+ x = x.view(n, int(h * scale_factor), int(w * scale_factor),
171
+ int(c / (scale_factor * scale_factor)))
172
+ return x
173
+
174
+ def extract_feature(self, pixel_values):
175
+ vit_embeds = self.vision_model(
176
+ pixel_values=pixel_values,
177
+ output_hidden_states=True,
178
+ return_dict=True).hidden_states[-4]
179
+ vit_embeds = vit_embeds[:, 1:, :]
180
+ # if torch.distributed.get_rank() == 0:
181
+ # print("before pixel shuffle:", vit_embeds.shape)
182
+ h = w = int(vit_embeds.shape[1] ** 0.5)
183
+ vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
184
+ vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=0.5)
185
+ vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
186
+ # if torch.distributed.get_rank() == 0:
187
+ # print("after pixel shuffle:", vit_embeds.shape)
188
+ vit_embeds = self.mlp1(vit_embeds)
189
+ return vit_embeds
190
+
191
+ def chat(self, tokenizer, pixel_values, question, generation_config,
192
+ IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'):
193
+
194
+ img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
195
+ self.img_context_token_id = img_context_token_id
196
+
197
+ from .conversation import get_conv_template
198
+
199
+ template = get_conv_template(self.template)
200
+ image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token + IMG_END_TOKEN
201
+ template.append_message(template.roles[0], image_tokens + '\n' + question)
202
+ template.append_message(template.roles[1], None)
203
+ query = template.get_prompt()
204
+ model_inputs = tokenizer(query, return_tensors='pt')
205
+ input_ids = model_inputs['input_ids'].cuda()
206
+ attention_mask = model_inputs['attention_mask'].cuda()
207
+
208
+ generation_output = self.generate(
209
+ pixel_values=pixel_values,
210
+ input_ids=input_ids,
211
+ attention_mask=attention_mask,
212
+ **generation_config
213
+ )
214
+ response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
215
+ query_to_print = query.replace(image_tokens, '<image>')
216
+ print(query_to_print, response)
217
+ return response
218
+
219
+ @torch.no_grad()
220
+ def generate(
221
+ self,
222
+ pixel_values: Optional[torch.FloatTensor] = None,
223
+ input_ids: Optional[torch.FloatTensor] = None,
224
+ attention_mask: Optional[torch.LongTensor] = None,
225
+ visual_features: Optional[torch.FloatTensor] = None,
226
+ generation_config: Optional[GenerationConfig] = None,
227
+ output_hidden_states: Optional[bool] = None,
228
+ return_dict: Optional[bool] = None,
229
+ **generate_kwargs,
230
+ ) -> torch.LongTensor:
231
+
232
+ assert self.img_context_token_id is not None
233
+ if pixel_values is not None:
234
+ if visual_features is not None:
235
+ vit_embeds = visual_features
236
+ else:
237
+ vit_embeds = self.extract_feature(pixel_values)
238
+
239
+ input_embeds = self.language_model.get_input_embeddings()(input_ids)
240
+ B, N, C = input_embeds.shape
241
+ input_embeds = input_embeds.reshape(B * N, C)
242
+
243
+ input_ids = input_ids.reshape(B * N)
244
+ selected = (input_ids == self.img_context_token_id)
245
+ assert selected.sum() != 0
246
+ input_embeds[selected] = vit_embeds.reshape(-1, C)
247
+
248
+ input_embeds = input_embeds.reshape(B, N, C)
249
+ else:
250
+ input_embeds = self.language_model.get_input_embeddings()(input_ids)
251
+
252
+ outputs = self.language_model.generate(
253
+ inputs_embeds=input_embeds,
254
+ attention_mask=attention_mask,
255
+ generation_config=generation_config,
256
+ output_hidden_states=output_hidden_states,
257
+ return_dict=return_dict,
258
+ use_cache=True,
259
+ **generate_kwargs,
260
+ )
261
+
262
+ return outputs