Refactor UI and fix some bugs (#1300)

Author: tastelikefeet

swift/llm/utils/client_utils.py

@@ -169,7 +169,7 @@ def _pre_inference_client(model_type: str,
         url = f'http://{host}:{port}/v1'
     url = url.rstrip('/')
     if is_chat_request:
-        messages = history_to_messages(history, query, system)
+        messages = history_to_messages(history, query, system, kwargs.get('roles'))
         if is_multimodal:
             messages = convert_to_base64(messages=messages)['messages']
             images = convert_to_base64(images=images)['images']

swift/llm/utils/utils.py

@@ -820,18 +820,22 @@ def compute_token_length(history_length: int) -> int:
 
 def history_to_messages(history: Optional[History],
                         query: Optional[str] = None,
-                        system: Optional[str] = None) -> Messages:
+                        system: Optional[str] = None,
+                        roles: Optional[List[List[str]]] = None) -> Messages:
     if history is None:
         history = []
     messages = []
+    if not roles:
+        roles = [['user', 'assistant']] * (len(history) + 1)
+    assert len(roles) == len(history) + 1
     if system is not None:
         messages.append({'role': 'system', 'content': system})
-    for h in history:
+    for role, h in zip(roles, history):
         assert isinstance(h, (list, tuple))
-        messages.append({'role': 'user', 'content': h[0]})
-        messages.append({'role': 'assistant', 'content': h[1]})
+        messages.append({'role': role[0], 'content': h[0]})
+        messages.append({'role': role[1], 'content': h[1]})
     if query is not None:
-        messages.append({'role': 'user', 'content': query})
+        messages.append({'role': roles[-1][0], 'content': query})
     return messages
 
 

swift/ui/llm_eval/eval.py

@@ -4,6 +4,9 @@
 import gradio as gr
 
 from swift.ui.base import BaseUI
+from swift.utils import get_logger
+
+logger = get_logger()
 
 
 class Eval(BaseUI):
@@ -110,10 +113,22 @@ class Eval(BaseUI):
 
     @classmethod
     def do_build_ui(cls, base_tab: Type['BaseUI']):
+        try:
+            from llmuses.backend.opencompass import OpenCompassBackendManager
+        except ImportError as e:
+            logger.error('You are using web-ui, please '
+                         'install requirements by `pip install llmuses ms-opencompass -U`')
+            raise e
+
         with gr.Row():
             gr.Textbox(elem_id='name', scale=20)
             gr.Dropdown(
-                elem_id='eval_dataset', is_list=True, choices=['ceval', 'gsm8k', 'arc'], multiselect=True, scale=20)
+                elem_id='eval_dataset',
+                is_list=True,
+                choices=OpenCompassBackendManager.list_datasets(),
+                multiselect=True,
+                allow_custom_value=True,
+                scale=20)
             gr.Textbox(elem_id='eval_few_shot', scale=20)
             gr.Textbox(elem_id='eval_limit', scale=20)
             gr.Checkbox(elem_id='eval_use_cache', scale=20)

swift/ui/llm_infer/llm_infer.py

@@ -13,8 +13,8 @@
 from gradio import Accordion, Tab
 from modelscope import GenerationConfig, snapshot_download
 
-from swift.llm import (DeployArguments, InferArguments, XRequestConfig, inference_client, inference_stream,
-                       limit_history_length, prepare_model_template)
+from swift.llm import (TEMPLATE_MAPPING, DeployArguments, InferArguments, XRequestConfig, inference_client,
+                       inference_stream, limit_history_length, prepare_model_template)
 from swift.ui.base import BaseUI
 from swift.ui.llm_infer.model import Model
 from swift.ui.llm_infer.runtime import Runtime
@@ -349,7 +349,7 @@ def prepare_checkpoint(cls, *args):
 
     @classmethod
     def clear_session(cls):
-        return '', [], None, []
+        return '', [], gr.update(value=None, interactive=True), []
 
     @classmethod
     def change_interactive(cls):
@@ -365,6 +365,14 @@ def _replace_tag_with_media(cls, history):
                 total_history.append(h[:2])
         return total_history
 
+    @classmethod
+    def agent_type(cls, response):
+        if response.lower().endswith('observation:'):
+            return 'react'
+        if 'observation:' not in response.lower() and 'action input:' in response.lower():
+            return 'toolbench'
+        return None
+
     @classmethod
     def send_message(cls, running_task, model_and_template, template_type, prompt: str, image, history, system,
                      max_new_tokens, temperature, top_k, top_p, repetition_penalty):
@@ -393,20 +401,38 @@ def send_message(cls, running_task, model_and_template, template_type, prompt: s
         request_config.stop = ['Observation:']
         stream_resp_with_history = ''
         medias = [m for h in old_history for m in h[2]]
+        media_infer_type = TEMPLATE_MAPPING[template].get('infer_media_type', 'round')
+        image_interactive = media_infer_type != 'dialogue'
+
+        text_history = [h for h in old_history if h[0]]
+        roles = []
+        for i in range(len(text_history) + 1):
+            roles.append(['user', 'assistant'])
+
+        for i, h in enumerate(text_history):
+            agent_type = cls.agent_type(h[1])
+            if i < len(text_history) - 1 and agent_type == 'toolbench':
+                roles[i + 1][0] = 'tool'
+            if i == len(text_history) - 1 and agent_type in ('toolbench', 'react'):
+                roles[i + 1][0] = 'tool'
+
         if not template_type.endswith('generation'):
             stream_resp = inference_client(
                 model_type,
                 prompt,
                 images=medias,
-                history=[h[:2] for h in old_history if h[0]],
+                history=[h[:2] for h in text_history],
                 system=system,
                 port=args['port'],
-                request_config=request_config)
+                request_config=request_config,
+                roles=roles,
+            )
             for chunk in stream_resp:
                 stream_resp_with_history += chunk.choices[0].delta.content
                 old_history[-1][0] = prompt
                 old_history[-1][1] = stream_resp_with_history
-                yield '', cls._replace_tag_with_media(old_history), None, old_history
+                yield ('', cls._replace_tag_with_media(old_history),
+                       gr.update(value=None, interactive=image_interactive), old_history)
         else:
             request_config.max_tokens = max_new_tokens
             stream_resp = inference_client(
@@ -415,7 +441,8 @@ def send_message(cls, running_task, model_and_template, template_type, prompt: s
                 stream_resp_with_history += chunk.choices[0].text
                 old_history[-1][0] = prompt
                 old_history[-1][1] = stream_resp_with_history
-                yield '', cls._replace_tag_with_media(old_history), None, old_history
+                yield ('', cls._replace_tag_with_media(old_history),
+                       gr.update(value=None, interactive=image_interactive), old_history)
 
     @classmethod
     def generate_chat(cls, model_and_template, template_type, prompt: str, image, history, system, max_new_tokens,

swift/ui/llm_train/advanced.py

@@ -76,6 +76,86 @@ class Advanced(BaseUI):
                 'en': 'Input in the json format'
             }
         },
+        'custom_train_dataset_path': {
+            'label': {
+                'zh': '自定义训练数据集路径',
+                'en': 'Custom train dataset path'
+            },
+            'info': {
+                'zh': '输入自定义的训练数据集路径，空格分隔',
+                'en': 'Extra train files, split by blank'
+            }
+        },
+        'custom_val_dataset_path': {
+            'label': {
+                'zh': '自定义校验数据集路径',
+                'en': 'Custom val dataset path'
+            },
+            'info': {
+                'zh': '输入自定义的校验数据集路径，逗号分隔',
+                'en': 'Extra val files, split by comma'
+            }
+        },
+        'truncation_strategy': {
+            'label': {
+                'zh': '数据集超长策略',
+                'en': 'Dataset truncation strategy'
+            },
+            'info': {
+                'zh': '如果token超长该如何处理',
+                'en': 'How to deal with the rows exceed the max length'
+            }
+        },
+        'gpu_memory_fraction': {
+            'label': {
+                'zh': 'GPU显存限制',
+                'en': 'GPU memory fraction'
+            },
+            'info': {
+                'zh': '设置使用显存的比例，一般用于显存测试',
+                'en': 'Set the memory fraction ratio of GPU, usually used in memory test'
+            }
+        },
+        'max_steps': {
+            'label': {
+                'zh': '最大迭代步数',
+                'en': 'Max steps',
+            },
+            'info': {
+                'zh': '设置最大迭代步数，该值如果大于零则数据集迭代次数不生效',
+                'en': 'Set the max steps, if the value > 0 then num_train_epochs has no effects',
+            }
+        },
+        'eval_batch_size': {
+            'label': {
+                'zh': '验证batch size',
+                'en': 'Val batch size',
+            },
+            'info': {
+                'zh': '验证的batch size',
+                'en': 'Set the val batch size',
+            }
+        },
+        'max_grad_norm': {
+            'label': {
+                'zh': '梯度裁剪',
+                'en': 'Max grad norm',
+            },
+            'info': {
+                'zh': '设置梯度裁剪',
+                'en': 'Set the max grad norm',
+            }
+        },
+        'predict_with_generate': {
+            'label': {
+                'zh': '使用生成指标代替loss',
+                'en': 'Use generate metric instead of loss',
+            },
+            'info': {
+                'zh': '验证时使用generate/Rouge代替loss',
+                'en': 'Use model.generate/Rouge instead of loss',
+            }
+        },
     }
 
     @classmethod
@@ -87,6 +167,16 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                     gr.Textbox(elem_id='weight_decay', lines=1, scale=20)
                     gr.Textbox(elem_id='logging_steps', lines=1, scale=20)
                     gr.Textbox(elem_id='lr_scheduler_type', lines=1, scale=20)
+                    gr.Textbox(elem_id='max_steps', lines=1, scale=20)
                     gr.Slider(elem_id='warmup_ratio', minimum=0.0, maximum=1.0, step=0.05, scale=20)
+                with gr.Row():
+                    gr.Textbox(elem_id='custom_train_dataset_path', is_list=True, scale=20)
+                    gr.Textbox(elem_id='custom_val_dataset_path', is_list=True, scale=20)
+                    gr.Dropdown(elem_id='truncation_strategy', scale=20)
+                    gr.Slider(elem_id='eval_batch_size', minimum=1, maximum=256, step=2, scale=20)
+                    gr.Textbox(elem_id='max_grad_norm', lines=1, scale=20)
+                    gr.Checkbox(elem_id='predict_with_generate', scale=20)
+                with gr.Row():
+                    gr.Textbox(elem_id='gpu_memory_fraction', scale=4)
                 with gr.Row():
                     gr.Textbox(elem_id='more_params', lines=4, scale=20)

swift/ui/llm_train/dataset.py

@@ -32,26 +32,6 @@ class Dataset(BaseUI):
                 'en': 'Set the max length input to the model',
             }
         },
-        'custom_train_dataset_path': {
-            'label': {
-                'zh': '自定义训练数据集路径',
-                'en': 'Custom train dataset path'
-            },
-            'info': {
-                'zh': '输入自定义的训练数据集路径，空格分隔',
-                'en': 'Extra train files, split by blank'
-            }
-        },
-        'custom_val_dataset_path': {
-            'label': {
-                'zh': '自定义校验数据集路径',
-                'en': 'Custom val dataset path'
-            },
-            'info': {
-                'zh': '输入自定义的校验数据集路径，逗号分隔',
-                'en': 'Extra val files, split by comma'
-            }
-        },
         'dataset_test_ratio': {
             'label': {
                 'zh': '验证集拆分比例',
@@ -82,16 +62,6 @@ class Dataset(BaseUI):
                 'en': 'Validate with the sample size from the dataset',
             }
         },
-        'truncation_strategy': {
-            'label': {
-                'zh': '数据集超长策略',
-                'en': 'Dataset truncation strategy'
-            },
-            'info': {
-                'zh': '如果token超长该如何处理',
-                'en': 'How to deal with the rows exceed the max length'
-            }
-        },
         'custom_dataset_info': {
             'label': {
                 'zh': '外部数据集配置',
@@ -102,18 +72,27 @@ class Dataset(BaseUI):
                 'en': 'An extra dataset config to register your own datasets'
             }
         },
+        'dataset_param': {
+            'label': {
+                'zh': '数据集设置',
+                'en': 'Dataset settings'
+            },
+        },
     }
 
     @classmethod
     def do_build_ui(cls, base_tab: Type['BaseUI']):
-        with gr.Row():
-            gr.Dropdown(elem_id='dataset', multiselect=True, choices=list(DATASET_MAPPING.keys()), scale=20)
-            gr.Textbox(elem_id='custom_dataset_info', is_list=False, scale=20)
-            gr.Textbox(elem_id='custom_train_dataset_path', is_list=True, scale=20)
-            gr.Textbox(elem_id='custom_val_dataset_path', is_list=True, scale=20)
-        with gr.Row():
-            gr.Slider(elem_id='dataset_test_ratio', minimum=0.0, maximum=1.0, step=0.05, scale=20)
-            gr.Slider(elem_id='max_length', minimum=32, maximum=32768, step=32, scale=20)
-            gr.Textbox(elem_id='train_dataset_sample', scale=20)
-            gr.Textbox(elem_id='val_dataset_sample', scale=20)
-            gr.Dropdown(elem_id='truncation_strategy', scale=20)
+        with gr.Accordion(elem_id='dataset_param', open=True):
+            with gr.Row():
+                gr.Dropdown(
+                    elem_id='dataset',
+                    multiselect=True,
+                    choices=list(DATASET_MAPPING.keys()),
+                    scale=20,
+                    allow_custom_value=True)
+                gr.Textbox(elem_id='custom_dataset_info', is_list=False, scale=20)
+            with gr.Row():
+                gr.Slider(elem_id='dataset_test_ratio', minimum=0.0, maximum=1.0, step=0.05, scale=20)
+                gr.Slider(elem_id='max_length', minimum=32, maximum=32768, step=1, scale=20)
+                gr.Textbox(elem_id='train_dataset_sample', scale=20)
+                gr.Textbox(elem_id='val_dataset_sample', scale=20)