lint

swift/llm/export.py

@@ -1,6 +1,6 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
-from typing import Optional
+from typing import List, Optional
 
 import json
 import torch
@@ -93,6 +93,25 @@ def gptq_model_quantize(model, tokenizer):
     return gptq_quantizer
 
 
+def replace_and_concat(template: 'Template', template_list: List, placeholder: str, keyword: str):
+    final_str = ''
+    for t in template_list:
+        if isinstance(t, str):
+            final_str += t.replace(placeholder, keyword)
+        elif isinstance(t, (tuple, list)):
+            if isinstance(t[0], int):
+                final_str += template.tokenizer.decode(t)
+            else:
+                for attr in t:
+                    if attr == 'bos_token_id':
+                        final_str += template.tokenizer.bos_token
+                    elif attr == 'eos_token_id':
+                        final_str += template.tokenizer.eos_token
+                    else:
+                        raise ValueError(f'Unknown token: {attr}')
+    return final_str
+
+
 def llm_export(args: ExportArguments) -> None:
     global _args, template
     logger.info(f'args: {args}')
@@ -131,14 +150,14 @@ def llm_export(args: ExportArguments) -> None:
         with open(os.path.join(args.ollama_output_dir, 'Modelfile'), 'w') as f:
             f.write(f'FROM {model_dir}\n')
             f.write(f'TEMPLATE """{{{{ if .System }}}}'
-                    f'{template.system_prefix[0].replace("{{SYSTEM}}", "{{ .System }}")}'
+                    f'{replace_and_concat(template, template.system_prefix, "{{SYSTEM}}", "{{ .System }}")}'
                     f'{{{{ end }}}}')
             f.write(f'{{{{ if .Prompt }}}}'
-                    f'{template.prompt[0].replace("{{QUERY}}", "{{ .Prompt }}")}'
+                    f'{replace_and_concat(template, template.prompt, "{{QUERY}}", "{{ .Prompt }}")}'
                     f'{{{{ end }}}}')
             f.write('{{ .Response }}')
             f.write(template.suffix[0] + '"""\n')
-            f.write(f'PARAMETER stop "{template.suffix[0]}"\n')
+            f.write(f'PARAMETER stop "{replace_and_concat(template, template.suffix, "", "")}"\n')
             if args.stop_words:
                 for stop_word in args.stop_words:
                     f.write(f'PARAMETER stop "{stop_word}"\n')

swift/llm/utils/vision_utils.py

@@ -161,10 +161,10 @@ def draw_plot(img_dir: str, bbox: List[int], bbox_type: str, output_file: str):
     Template.normalize_bbox(objects, [image], 'real')
     bbox = objects[0]['bbox']
     draw = ImageDraw.Draw(image)
-    draw.rectangle(bbox, outline="red", width=2)
+    draw.rectangle(bbox, outline='red', width=2)
     image.save(output_file)
 
 
 if __name__ == '__main__':
     # A test main to draw bbox
-    draw_plot('/mnt/workspace/man.jpg', [354, 462, 580, 738], 'norm_1000', '/mnt/workspace/man_bbox.jpg')
+    draw_plot('/mnt/workspace/man.jpg', [354, 462, 580, 738], 'norm_1000', '/mnt/workspace/man_bbox.jpg')

tests/llm/test_ollama_export.py

@@ -35,7 +35,6 @@ def test_llama3(self):
             self.assertTrue(template in content)
             self.assertTrue(stop in content)
 
-    @unittest.skip('TODO FIX')
     def test_glm4(self):
         args = ExportArguments(model_type='glm4-9b-chat', to_ollama=True, ollama_output_dir=self.tmp_dir)
         export_main(args)