wehub-resource-sync
81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
import os
|
|
import shutil
|
|
import tempfile
|
|
import transformers
|
|
import unittest
|
|
from packaging import version
|
|
|
|
from swift import ExportArguments, export_main
|
|
|
|
if __name__ == '__main__':
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
|
|
|
|
|
|
class TestTemplate(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
|
|
self.tmp_dir = tempfile.TemporaryDirectory().name
|
|
|
|
def tearDown(self):
|
|
if os.path.exists(self.tmp_dir):
|
|
shutil.rmtree(self.tmp_dir)
|
|
super().tearDown()
|
|
|
|
@unittest.skip('swift2.0')
|
|
def test_llama3(self):
|
|
args = ExportArguments(model_type='llama3-8b-instruct', to_ollama=True, ollama_output_dir=self.tmp_dir)
|
|
export_main(args)
|
|
|
|
template = ('TEMPLATE """{{ if .System }}<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n'
|
|
'{{ .System }}<|eot_id|>{{ else }}<|begin_of_text|>{{ end }}{{ if .Prompt }}<|start_header_id|>user'
|
|
'<|end_header_id|>\n\n{{ .Prompt }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n'
|
|
'{{ end }}{{ .Response }}<|eot_id|>"""')
|
|
|
|
stop = 'PARAMETER stop "<|eot_id|>"'
|
|
|
|
with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
|
|
content = f.read()
|
|
self.assertTrue(template in content)
|
|
self.assertTrue(stop in content)
|
|
|
|
@unittest.skip('swift2.0')
|
|
def test_chatglm4(self):
|
|
if version.parse(transformers.__version__) >= version.parse('4.45'):
|
|
return
|
|
|
|
args = ExportArguments(model_type='glm4-9b-chat', to_ollama=True, ollama_output_dir=self.tmp_dir)
|
|
export_main(args)
|
|
|
|
template = ('TEMPLATE """{{ if .System }}[gMASK] <sop><|system|>\n{{ .System }}{{ else }}'
|
|
'[gMASK] <sop>{{ end }}{{ if .Prompt }}<|user|>\n{{ .Prompt }}<|assistant|>\n'
|
|
'{{ end }}{{ .Response }}<|user|>"""')
|
|
|
|
stop = 'PARAMETER stop "<|user|>"'
|
|
|
|
with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
|
|
content = f.read()
|
|
self.assertTrue(template in content)
|
|
self.assertTrue(stop in content)
|
|
|
|
@unittest.skip('swift2.0')
|
|
def test_qwen2(self):
|
|
args = ExportArguments(model_type='qwen2-7b-instruct', to_ollama=True, ollama_output_dir=self.tmp_dir)
|
|
export_main(args)
|
|
|
|
template = ('TEMPLATE """{{ if .System }}<|im_start|>system\n{{ .System }}<|im_end|>\n{{ else }}{{ end }}'
|
|
'{{ if .Prompt }}<|im_start|>user\n{{ .Prompt }}<|im_end|>\n<|im_start|>assistant\n'
|
|
'{{ end }}{{ .Response }}<|im_end|>"""')
|
|
|
|
stop = 'PARAMETER stop "<|im_end|>"'
|
|
|
|
with open(os.path.join(self.tmp_dir, 'Modelfile'), 'r') as f:
|
|
content = f.read()
|
|
self.assertTrue(template in content)
|
|
self.assertTrue(stop in content)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|