File size: 3,298 Bytes
cb2428f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from typing import List
from swift.llm import ExportArguments, PtEngine, RequestConfig, Template, prepare_model_template
from swift.utils import get_logger
logger = get_logger()
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 export_to_ollama(args: ExportArguments):
args.device_map = 'meta' # Accelerate load speed.
logger.info('Exporting to ollama:')
os.makedirs(args.output_dir, exist_ok=True)
model, template = prepare_model_template(args)
pt_engine = PtEngine.from_model_template(model, template)
logger.info(f'Using model_dir: {pt_engine.model_dir}')
template_meta = template.template_meta
with open(os.path.join(args.output_dir, 'Modelfile'), 'w', encoding='utf-8') as f:
f.write(f'FROM {pt_engine.model_dir}\n')
f.write(f'TEMPLATE """{{{{ if .System }}}}'
f'{replace_and_concat(template, template_meta.system_prefix, "{{SYSTEM}}", "{{ .System }}")}'
f'{{{{ else }}}}{replace_and_concat(template, template_meta.prefix, "", "")}'
f'{{{{ end }}}}')
f.write(f'{{{{ if .Prompt }}}}'
f'{replace_and_concat(template, template_meta.prompt, "{{QUERY}}", "{{ .Prompt }}")}'
f'{{{{ end }}}}')
f.write('{{ .Response }}')
f.write(replace_and_concat(template, template_meta.suffix, '', '') + '"""\n')
f.write(f'PARAMETER stop "{replace_and_concat(template, template_meta.suffix, "", "")}"\n')
request_config = RequestConfig(
temperature=args.temperature,
top_k=args.top_k,
top_p=args.top_p,
repetition_penalty=args.repetition_penalty)
generation_config = pt_engine._prepare_generation_config(request_config)
pt_engine._add_stop_words(generation_config, request_config, template.template_meta)
for stop_word in generation_config.stop_words:
f.write(f'PARAMETER stop "{stop_word}"\n')
f.write(f'PARAMETER temperature {generation_config.temperature}\n')
f.write(f'PARAMETER top_k {generation_config.top_k}\n')
f.write(f'PARAMETER top_p {generation_config.top_p}\n')
f.write(f'PARAMETER repeat_penalty {generation_config.repetition_penalty}\n')
logger.info('Save Modelfile done, you can start ollama by:')
logger.info('> ollama serve')
logger.info('In another terminal:')
logger.info('> ollama create my-custom-model ' f'-f {os.path.join(args.output_dir, "Modelfile")}')
logger.info('> ollama run my-custom-model')
|