# Copyright (c) Alibaba, Inc. and its affiliates. import os import re import sys import time from datetime import datetime from functools import partial from typing import Type import gradio as gr import json import torch from json import JSONDecodeError from transformers.utils import is_torch_cuda_available, is_torch_npu_available from swift.llm import ExportArguments from swift.ui.base import BaseUI from swift.ui.llm_export.export import Export from swift.ui.llm_export.model import Model from swift.ui.llm_export.runtime import ExportRuntime from swift.utils import get_device_count class LLMExport(BaseUI): group = 'llm_export' sub_ui = [Model, Export, ExportRuntime] locale_dict = { 'llm_export': { 'label': { 'zh': 'LLM导出', 'en': 'LLM export', } }, 'more_params': { 'label': { 'zh': '更多参数', 'en': 'More params' }, 'info': { 'zh': '以json格式或--xxx xxx命令行格式填入', 'en': 'Fill in with json format or --xxx xxx cmd format' } }, 'export': { 'value': { 'zh': '开始导出', 'en': 'Begin Export' }, }, 'gpu_id': { 'label': { 'zh': '选择可用GPU', 'en': 'Choose GPU' }, 'info': { 'zh': '选择使用的GPU号，如CUDA不可用只能选择CPU', 'en': 'Select GPU to export' } }, } choice_dict = BaseUI.get_choices_from_dataclass(ExportArguments) default_dict = BaseUI.get_default_value_from_dataclass(ExportArguments) arguments = BaseUI.get_argument_names(ExportArguments) @classmethod def do_build_ui(cls, base_tab: Type['BaseUI']): with gr.TabItem(elem_id='llm_export', label=''): default_device = 'cpu' device_count = get_device_count() if device_count > 0: default_device = '0' with gr.Blocks(): Model.build_ui(base_tab) Export.build_ui(base_tab) ExportRuntime.build_ui(base_tab) with gr.Row(): gr.Textbox(elem_id='more_params', lines=4, scale=20) gr.Button(elem_id='export', scale=2, variant='primary') gr.Dropdown( elem_id='gpu_id', multiselect=True, choices=[str(i) for i in range(device_count)] + ['cpu'], value=default_device, scale=8) cls.element('export').click( cls.export_model, list(base_tab.valid_elements().values()), [cls.element('runtime_tab'), cls.element('running_tasks')]) base_tab.element('running_tasks').change( partial(ExportRuntime.task_changed, base_tab=base_tab), [base_tab.element('running_tasks')], list(base_tab.valid_elements().values()) + [cls.element('log')]) ExportRuntime.element('kill_task').click( ExportRuntime.kill_task, [ExportRuntime.element('running_tasks')], [ExportRuntime.element('running_tasks')] + [ExportRuntime.element('log')], ) @classmethod def export(cls, *args): export_args = cls.get_default_value_from_dataclass(ExportArguments) kwargs = {} kwargs_is_list = {} other_kwargs = {} more_params = {} more_params_cmd = '' keys = cls.valid_element_keys() for key, value in zip(keys, args): compare_value = export_args.get(key) compare_value_arg = str(compare_value) if not isinstance(compare_value, (list, dict)) else compare_value compare_value_ui = str(value) if not isinstance(value, (list, dict)) else value if key in export_args and compare_value_ui != compare_value_arg and value: if isinstance(value, str) and re.fullmatch(cls.int_regex, value): value = int(value) elif isinstance(value, str) and re.fullmatch(cls.float_regex, value): value = float(value) elif isinstance(value, str) and re.fullmatch(cls.bool_regex, value): value = True if value.lower() == 'true' else False kwargs[key] = value if not isinstance(value, list) else ' '.join(value) kwargs_is_list[key] = isinstance(value, list) or getattr(cls.element(key), 'is_list', False) else: other_kwargs[key] = value if key == 'more_params' and value: try: more_params = json.loads(value) except (JSONDecodeError or TypeError): more_params_cmd = value kwargs.update(more_params) model = kwargs.get('model') if os.path.exists(model) and os.path.exists(os.path.join(model, 'args.json')): kwargs['ckpt_dir'] = kwargs.pop('model') export_args = ExportArguments( **{ key: value.split(' ') if key in kwargs_is_list and kwargs_is_list[key] else value for key, value in kwargs.items() }) params = '' sep = f'{cls.quote} {cls.quote}' for e in kwargs: if isinstance(kwargs[e], list): params += f'--{e} {cls.quote}{sep.join(kwargs[e])}{cls.quote} ' elif e in kwargs_is_list and kwargs_is_list[e]: all_args = [arg for arg in kwargs[e].split(' ') if arg.strip()] params += f'--{e} {cls.quote}{sep.join(all_args)}{cls.quote} ' else: params += f'--{e} {cls.quote}{kwargs[e]}{cls.quote} ' params += more_params_cmd + ' ' devices = other_kwargs['gpu_id'] devices = [d for d in devices if d] assert (len(devices) == 1 or 'cpu' not in devices) gpus = ','.join(devices) cuda_param = '' if gpus != 'cpu': if is_torch_npu_available(): cuda_param = f'ASCEND_RT_VISIBLE_DEVICES={gpus}' elif is_torch_cuda_available(): cuda_param = f'CUDA_VISIBLE_DEVICES={gpus}' else: cuda_param = '' now = datetime.now() time_str = f'{now.year}{now.month}{now.day}{now.hour}{now.minute}{now.second}' file_path = f'output/{export_args.model_type}-{time_str}' if not os.path.exists(file_path): os.makedirs(file_path, exist_ok=True) log_file = os.path.join(os.getcwd(), f'{file_path}/run_export.log') export_args.log_file = log_file params += f'--log_file "{log_file}" ' params += '--ignore_args_error true ' additional_param = '' if export_args.quant_method == 'gptq': additional_param = 'OMP_NUM_THREADS=14' if sys.platform == 'win32': if cuda_param: cuda_param = f'set {cuda_param} && ' if additional_param: additional_param = f'set {additional_param} && ' run_command = f'{cuda_param}{additional_param}start /b swift export {params} > {log_file} 2>&1' else: run_command = f'{cuda_param} {additional_param} nohup swift export {params} > {log_file} 2>&1 &' return run_command, export_args, log_file @classmethod def export_model(cls, *args): run_command, export_args, log_file = cls.export(*args) os.system(run_command) time.sleep(2) return gr.update(open=True), ExportRuntime.refresh_tasks(log_file)