base.py

# Copyright (c) ModelScope Contributors. All rights reserved.
import dataclasses
import gradio as gr
import json
import os
import sys
import time
import typing
from collections import OrderedDict
from dataclasses import fields
from datetime import datetime
from functools import wraps
from gradio import Accordion, Audio, Button, Checkbox, Dropdown, File, Image, Slider, Tab, TabItem, Textbox, Video
from modelscope.hub.utils.utils import get_cache_dir
from typing import Any, Dict, List, Literal, Optional, Type, Union, get_args, get_origin

from swift.arguments import BaseArguments
from swift.model import get_matched_model_meta
from swift.template import TEMPLATE_MAPPING

all_langs = ['zh', 'en']
builder: Type['BaseUI'] = None
base_builder: Type['BaseUI'] = None

DEFAULT_GRPO_SYSTEM = (
    'A conversation between User and Assistant. The user asks a question, and the Assistant solves it. '
    'The assistant first thinks about the reasoning process in the mind and then provides the user '
    'with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> '
    '</answer> tags, respectively, i.e., <think> reasoning process here </think>'
    '<answer> answer here </answer>')


def update_data(fn):

    @wraps(fn)
    def wrapper(*args, **kwargs):
        elem_id = kwargs.get('elem_id', None)
        self = args[0]

        if builder is not None:
            choices = base_builder.choice(elem_id)
            if choices:
                choices = [str(choice) if choice is not None else None for choice in choices]
                kwargs['choices'] = choices

        if not isinstance(self, (Tab, TabItem, Accordion)) and 'interactive' not in kwargs:  # noqa
            kwargs['interactive'] = True

        if 'is_list' in kwargs:
            self.is_list = kwargs.pop('is_list')

        if base_builder and base_builder.default(elem_id) is not None and not kwargs.get('value'):
            kwargs['value'] = base_builder.default(elem_id)

        if builder is not None:
            if elem_id in builder.locales(builder.lang):
                values = builder.locale(elem_id, builder.lang)
                if 'info' in values:
                    kwargs['info'] = values['info']
                if 'value' in values:
                    kwargs['value'] = values['value']
                if 'label' in values:
                    kwargs['label'] = values['label']
                if hasattr(builder, 'visible'):
                    kwargs['visible'] = builder.visible
                argument = base_builder.argument(elem_id)
                if argument and 'label' in kwargs:
                    kwargs['label'] = kwargs['label'] + f'({argument})'

        kwargs['elem_classes'] = 'align'
        ret = fn(self, **kwargs)
        self.constructor_args.update(kwargs)

        if builder is not None:
            builder.element_dict[elem_id] = self
        return ret

    return wrapper


Textbox.__init__ = update_data(Textbox.__init__)
Dropdown.__init__ = update_data(Dropdown.__init__)
Checkbox.__init__ = update_data(Checkbox.__init__)
Slider.__init__ = update_data(Slider.__init__)
TabItem.__init__ = update_data(TabItem.__init__)
Accordion.__init__ = update_data(Accordion.__init__)
Button.__init__ = update_data(Button.__init__)
File.__init__ = update_data(File.__init__)
Image.__init__ = update_data(Image.__init__)
Video.__init__ = update_data(Video.__init__)
Audio.__init__ = update_data(Audio.__init__)


class BaseUI:

    choice_dict: Dict[str, List] = {}
    default_dict: Dict[str, Any] = {}
    locale_dict: Dict[str, Dict] = {}
    element_dict: Dict[str, Dict] = {}
    arguments: Dict[str, str] = {}
    sub_ui: List[Type['BaseUI']] = []
    group: str = None
    lang: str = all_langs[0]
    int_regex = r'^[-+]?[0-9]+$'
    float_regex = r'[-+]?(?:\d*\.*\d+)'
    bool_regex = r'^(T|t)rue$|^(F|f)alse$'
    cache_dir = os.path.join(get_cache_dir(), 'swift-web-ui')
    os.makedirs(cache_dir, exist_ok=True)
    quote = '\'' if sys.platform != 'win32' else '"'
    visible = True
    _locale = {
        'local_dir_alert': {
            'value': {
                'zh': '无法识别model_type和template,请手动选择',
                'en': 'Cannot recognize the model_type and template, please choose manually'
            }
        },
    }

    @classmethod
    def build_ui(cls, base_tab: Type['BaseUI']):
        """Build UI"""
        global builder, base_builder
        cls.element_dict = {}
        old_builder = builder
        old_base_builder = base_builder
        builder = cls
        base_builder = base_tab
        cls.do_build_ui(base_tab)
        builder = old_builder
        base_builder = old_base_builder
        if cls is base_tab:
            for ui in cls.sub_ui:
                ui.after_build_ui(base_tab)

    @classmethod
    def after_build_ui(cls, base_tab: Type['BaseUI']):
        pass

    @classmethod
    def do_build_ui(cls, base_tab: Type['BaseUI']):
        """Build UI"""
        pass

    @classmethod
    def save_cache(cls, key, value):
        timestamp = str(int(time.time()))
        key = key.replace('/', '-')
        filename = os.path.join(cls.cache_dir, key + '-' + timestamp)
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(value, f)

    @classmethod
    def list_cache(cls, key):
        files = []
        key = key.replace('/', '-')
        for _, _, filenames in os.walk(cls.cache_dir):
            for filename in filenames:
                if filename.startswith(key):
                    idx = filename.rfind('-')
                    key, ts = filename[:idx], filename[idx + 1:]
                    dt_object = datetime.fromtimestamp(int(ts))
                    formatted_time = dt_object.strftime('%Y/%m/%d %H:%M:%S')
                    files.append(formatted_time)
        return sorted(files, reverse=True)

    @classmethod
    def load_cache(cls, key, timestamp) -> BaseArguments:
        dt_object = datetime.strptime(timestamp, '%Y/%m/%d %H:%M:%S')
        timestamp = int(dt_object.timestamp())
        key = key.replace('/', '-')
        filename = key + '-' + str(timestamp)
        with open(os.path.join(cls.cache_dir, filename), 'r', encoding='utf-8') as f:
            return json.load(f)

    @classmethod
    def clear_cache(cls, key):
        key = key.replace('/', '-')
        for _, _, filenames in os.walk(cls.cache_dir):
            for filename in filenames:
                if filename.startswith(key):
                    os.remove(os.path.join(cls.cache_dir, filename))

    @classmethod
    def choice(cls, elem_id):
        """Get choice by elem_id"""
        for sub_ui in BaseUI.sub_ui:
            _choice = sub_ui.choice(elem_id)
            if _choice:
                return _choice
        return cls.choice_dict.get(elem_id, [])

    @classmethod
    def default(cls, elem_id):
        """Get choice by elem_id"""
        if elem_id in cls.default_dict:
            return cls.default_dict.get(elem_id)
        for sub_ui in BaseUI.sub_ui:
            _choice = sub_ui.default(elem_id)
            if _choice:
                return _choice
        return None

    @classmethod
    def locale(cls, elem_id, lang):
        """Get locale by elem_id"""
        return cls.locales(lang)[elem_id]

    @classmethod
    def locales(cls, lang):
        """Get locale by lang"""
        locales = OrderedDict()
        for sub_ui in cls.sub_ui:
            _locales = sub_ui.locales(lang)
            locales.update(_locales)
        for key, value in cls.locale_dict.items():
            locales[key] = {k: v[lang] for k, v in value.items()}
        return locales

    @classmethod
    def elements(cls):
        """Get all elements"""
        elements = OrderedDict()
        elements.update(cls.element_dict)
        for sub_ui in cls.sub_ui:
            _elements = sub_ui.elements()
            elements.update(_elements)
        return elements

    @classmethod
    def valid_elements(cls):
        valid_elements = OrderedDict()
        elements = cls.elements()
        for key, value in elements.items():
            if isinstance(value, (Textbox, Dropdown, Slider, Checkbox)) and key != 'train_record':
                valid_elements[key] = value
        return valid_elements

    @classmethod
    def element_keys(cls):
        return list(cls.elements().keys())

    @classmethod
    def valid_element_keys(cls):
        return [
            key for key, value in cls.elements().items()
            if isinstance(value, (Textbox, Dropdown, Slider, Checkbox)) and key != 'train_record'
        ]

    @classmethod
    def element(cls, elem_id):
        """Get element by elem_id"""
        elements = cls.elements()
        return elements[elem_id]

    @classmethod
    def argument(cls, elem_id):
        """Get argument by elem_id"""
        return cls.arguments.get(elem_id)

    @classmethod
    def set_lang(cls, lang):
        cls.lang = lang
        for sub_ui in cls.sub_ui:
            sub_ui.lang = lang

    @staticmethod
    def get_choices_from_dataclass(dataclass):
        choice_dict = {}
        for f in fields(dataclass):
            default_value = f.default
            type_orign = get_origin(f.type)
            type_args = get_args(f.type)
            if 'MISSING_TYPE' in str(default_value):
                default_value = None
            if 'choices' in f.metadata:
                choice_dict[f.name] = list(f.metadata['choices'])
            if type_orign is Literal:
                choice_dict[f.name] = list(type_args)
            elif type_orign is Union and type(None) in type_args:
                for inner_type in type_args:
                    if get_origin(inner_type) is Literal:
                        choice_dict[f.name] = list(get_args(inner_type))
                        break
            if f.name in choice_dict and default_value not in choice_dict[f.name]:
                choice_dict[f.name].insert(0, default_value)
        return choice_dict

    @staticmethod
    def get_default_value_from_dataclass(dataclass):
        default_dict = {}
        for f in fields(dataclass):
            if f.default.__class__ is dataclasses._MISSING_TYPE:
                default_dict[f.name] = f.default_factory()
            else:
                default_dict[f.name] = f.default
            if isinstance(default_dict[f.name], list):
                try:
                    default_dict[f.name] = ' '.join(default_dict[f.name])
                except TypeError:
                    default_dict[f.name] = None
            if not default_dict[f.name] and default_dict[f.name] not in (0, False):
                default_dict[f.name] = None
        return default_dict

    @staticmethod
    def get_argument_names(dataclass):
        arguments = {}
        for f in fields(dataclass):
            arguments[f.name] = f'--{f.name}'
        return arguments

    @classmethod
    def update_input_model(cls,
                           model,
                           allow_keys=None,
                           has_record=True,
                           arg_cls=BaseArguments,
                           is_ref_model=False,
                           is_reward_model=False):
        keys = cls.valid_element_keys()
        if allow_keys:
            keys = [key for key in keys if key in allow_keys]

        if not model:
            ret = [gr.update()] * (len(keys) + int(has_record))
            if len(ret) == 1:
                return ret[0]
            else:
                return ret

        model_meta = get_matched_model_meta(model)
        local_args_path = os.path.join(model, 'args.json')
        if model_meta is None and not os.path.exists(local_args_path):
            gr.Info(cls._locale['local_dir_alert']['value'][cls.lang])
            ret = [gr.update()] * (len(keys) + int(has_record))
            if len(ret) == 1:
                return ret[0]
            else:
                return ret

        if os.path.exists(local_args_path):
            try:
                if hasattr(arg_cls, 'resume_from_checkpoint'):
                    # Use adapter vs full-model constructor by path type (avoids fragile exception message check).
                    if BaseArguments._check_is_adapter(model):
                        args = arg_cls(resume_from_checkpoint=model, load_data_args=True)
                    else:
                        args = arg_cls(model=model, load_data_args=True)
                else:
                    if os.path.exists(os.path.join(model, 'adapter_config.json')):
                        args = arg_cls(adapters=model, load_data_args=True)
                    else:
                        args = arg_cls(model=model, load_data_args=True)
            except ValueError:
                return [gr.update()] * (len(keys) + int(has_record))
            values = []
            for key in keys:
                arg_value = getattr(args, key, None)
                if arg_value and key != 'model':
                    if key in ('torch_dtype', 'bnb_4bit_compute_dtype'):
                        arg_value = str(arg_value).split('.')[1]
                    if isinstance(arg_value, list) and key != 'dataset':
                        try:
                            arg_value = ' '.join(arg_value)
                        except Exception:
                            arg_value = None
                    values.append(gr.update(value=arg_value))
                else:
                    values.append(gr.update())
            ret = [gr.update(choices=[])] * int(has_record) + values
            if len(ret) == 1:
                return ret[0]
            else:
                return ret
        else:
            values = []
            for key in keys:
                if key not in ('template', 'model_type', 'ref_model_type', 'reward_model_type', 'system'):
                    values.append(gr.update())
                elif key in ('template', 'model_type', 'ref_model_type', 'reward_model_type'):
                    if key == 'ref_model_type':
                        if is_ref_model:
                            values.append(gr.update(value=getattr(model_meta, 'model_type')))
                        else:
                            values.append(gr.update())
                    elif key == 'reward_model_type':
                        if is_reward_model:
                            values.append(gr.update(value=getattr(model_meta, 'model_type')))
                        else:
                            values.append(gr.update())
                    else:
                        values.append(gr.update(value=getattr(model_meta, key)))
                else:
                    if cls.group == 'llm_grpo':
                        values.append(gr.update(value=DEFAULT_GRPO_SYSTEM))
                    else:
                        values.append(gr.update(value=TEMPLATE_MAPPING[model_meta.template].default_system))

        if has_record:
            return [gr.update(choices=cls.list_cache(model))] + values
        else:
            if len(values) == 1:
                return values[0]
            return values

    @classmethod
    def update_all_settings(cls, model, train_record, base_tab):
        if not train_record:
            return [gr.update()] * len(cls.valid_elements())
        cache = cls.load_cache(model, train_record)
        updates = []
        for key, value in base_tab.valid_elements().items():
            if key in cache:
                updates.append(gr.update(value=cache[key]))
            else:
                updates.append(gr.update())
        return updates

    @classmethod
    def update_ddp_num(cls, gpu_ids, use_ddp):
        if use_ddp:
            if 'cpu' in gpu_ids:
                return None
            else:
                return len(gpu_ids)
        return 1