Fix ui (#335)

Author: tastelikefeet

swift/ui/base.py

@@ -76,6 +76,8 @@ class BaseUI:
     sub_ui: List[Type['BaseUI']] = []
     group: str = None
     lang: str = all_langs[0]
+    int_regex = r'^[-+]?[0-9]+$'
+    float_regex = r'[-+]?(?:\d*\.*\d+)'
 
     @classmethod
     def build_ui(cls, base_tab: Type['BaseUI']):

swift/ui/llm_infer/llm_infer.py

@@ -7,6 +7,7 @@
 import torch
 from gradio import Accordion, Tab
 
+from swift import snapshot_download
 from swift.llm import (InferArguments, inference_stream, limit_history_length,
                        prepare_model_template)
 from swift.ui.base import BaseUI
@@ -19,9 +20,6 @@ class LLMInfer(BaseUI):
 
     sub_ui = [Model]
 
-    int_regex = r'^[-+]?[0-9]+$'
-    float_regex = r'[-+]?(?:\d*\.*\d+)'
-
     locale_dict = {
         'generate_alert': {
             'value': {
@@ -137,7 +135,6 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
 
     @classmethod
     def reset_load_button(cls):
-        gr.Info(cls.locale('loaded_alert', cls.lang)['value'])
         return gr.update(
             value=cls.locale('load_checkpoint', cls.lang)['value'])
 
@@ -184,7 +181,10 @@ def prepare_checkpoint(cls, *args):
 
         kwargs.update(more_params)
         if kwargs['model_type'] == cls.locale('checkpoint', cls.lang)['value']:
-            kwargs['ckpt_dir'] = kwargs.pop('model_id_or_path')
+            model_dir = kwargs.pop('model_id_or_path')
+            if not os.path.exists(model_dir):
+                model_dir = snapshot_download(model_dir)
+            kwargs['ckpt_dir'] = model_dir
         if 'ckpt_dir' in kwargs or 'model_id_or_path' in kwargs:
             kwargs.pop('model_type', None)
 
@@ -194,8 +194,9 @@ def prepare_checkpoint(cls, *args):
         gpus = ','.join(devices)
         if gpus != 'cpu':
             os.environ['CUDA_VISIBLE_DEVICES'] = gpus
-        inter_args = InferArguments(**kwargs)
-        model, template = prepare_model_template(inter_args)
+        infer_args = InferArguments(**kwargs)
+        model, template = prepare_model_template(infer_args)
+        gr.Info(cls.locale('loaded_alert', cls.lang)['value'])
         return [model, template]
 
     @classmethod

swift/ui/llm_infer/model.py

@@ -121,7 +121,7 @@ def update_input_model(choice):
                                                       cls.lang)['value'])
 
         def update_model_id_or_path(model_type, path):
-            if not path:
+            if not path or not os.path.exists(path):
                 return None, None, None
             local_path = os.path.join(path, 'sft_args.json')
             if not os.path.exists(local_path):

swift/ui/llm_train/llm_train.py

@@ -1,5 +1,6 @@
 import collections
 import os
+import re
 import sys
 import time
 from subprocess import PIPE, STDOUT, Popen
@@ -243,6 +244,13 @@ def train(cls, *args):
                 compare_value, (list, dict)) else compare_value
             compare_value_ui = str(value) if not isinstance(
                 value, (list, dict)) else 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)
+
             if key not in ignore_elements and key in sft_args and compare_value_ui != compare_value_arg and value:
                 kwargs[key] = value if not isinstance(
                     value, list) else ' '.join(value)