Upload modeling_skywork_chat.py

modeling_skywork_chat.py

@@ -1,4 +1,5 @@
 import warnings
+import re
 from typing import List, Optional, Tuple, Union
 
 import torch.utils.checkpoint
@@ -251,7 +252,7 @@ class SkyworkChatModel(PreTrainedModel):
 
     def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,
              num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
-             verbose=False):
+             verbose=False, mode="think"):
 
         if history is None and pixel_values is not None and '<image>' not in question:
             question = '<image>\n' + question
@@ -275,6 +276,8 @@ class SkyworkChatModel(PreTrainedModel):
         template.append_message(template.roles[0], question)
         template.append_message(template.roles[1], None)
         query = template.get_prompt()
+        if mode != "think":
+            query = re.sub(r'\n<think>', '', query, count=1)
    
 
         if verbose and pixel_values is not None: