| 241 | return outputs |
| 242 | |
| 243 | class Chat: |
| 244 | def __init__( |
| 245 | self, |
| 246 | model_path, |
| 247 | device, |
| 248 | num_gpus=1, |
| 249 | load_8bit=False, |
| 250 | temperature=0.7, |
| 251 | max_new_tokens=512, |
| 252 | lora_path=None, |
| 253 | ): |
| 254 | model, tokenizer = load_model( |
| 255 | model_path, device, num_gpus, load_8bit=load_8bit, lora_weights=lora_path |
| 256 | ) |
| 257 | |
| 258 | self.model = model |
| 259 | # self.model.language_model = deepspeed.init_inference( |
| 260 | # self.model.language_model, mp_size=1, dtype=torch.float16, checkpoint=None, replace_with_kernel_inject=True) |
| 261 | self.tokenizer = tokenizer |
| 262 | num_queries = model.config.num_query_tokens |
| 263 | |
| 264 | self.device = device |
| 265 | self.dtype = model.dtype |
| 266 | |
| 267 | stop_words = ["Human: ", "Assistant: ", "###", "\n\n"] |
| 268 | stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words] |
| 269 | stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)]) |
| 270 | |
| 271 | self.conv = get_conv_template("husky") |
| 272 | |
| 273 | self.image_query = DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN |
| 274 | self.video_query = DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN |
| 275 | |
| 276 | self.generation_config = GenerationConfig( |
| 277 | bos_token_id=1, |
| 278 | pad_token_id=0, |
| 279 | do_sample=True, |
| 280 | top_k=20, |
| 281 | top_p=0.9, |
| 282 | temperature=temperature, |
| 283 | max_new_tokens=max_new_tokens, |
| 284 | stopping_criteria=stopping_criteria |
| 285 | ) |
| 286 | |
| 287 | def ask(self, text, conv, modal_type="image"): |
| 288 | assert modal_type in ["text", "image", "video"] |
| 289 | conversations = [] |
| 290 | |
| 291 | if len(conv.messages) > 0 or modal_type == "text": |
| 292 | conv.append_message(conv.roles[0], text) |
| 293 | elif modal_type == "image": |
| 294 | conv.append_message(conv.roles[0], self.image_query + "\n" + text) |
| 295 | else: |
| 296 | conv.append_message(conv.roles[0], self.video_query + "\n" + text) |
| 297 | |
| 298 | conv.append_message(conv.roles[1], None) |
| 299 | conversations.append(conv.get_prompt()) |
| 300 | return conversations |