add yi-vl 6b model

README.md

@@ -166,6 +166,7 @@ Compared to ChatGLM's [P-Tuning](https://github.com/THUDM/ChatGLM2-6B/tree/main/
 | [StarCoder2](https://huggingface.co/bigcode)             | 3B/7B/15B                        | q_proj,v_proj     | -         |
 | [XVERSE](https://huggingface.co/xverse)                  | 7B/13B/65B                       | q_proj,v_proj     | xverse    |
 | [Yi (1/1.5)](https://huggingface.co/01-ai)               | 6B/9B/34B                        | q_proj,v_proj     | yi        |
+| [Yi-VL](https://huggingface.co/01-ai)                    | 6B                               | q_proj,v_proj     | yi_vl     |
 | [Yuan](https://huggingface.co/IEITYuan)                  | 2B/51B/102B                      | q_proj,v_proj     | yuan      |
 
 > [!NOTE]

README_zh.md

@@ -166,6 +166,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/ec36a9dd-37f4-4f72-81bd
 | [StarCoder2](https://huggingface.co/bigcode)             | 3B/7B/15B                        | q_proj,v_proj     | -         |
 | [XVERSE](https://huggingface.co/xverse)                  | 7B/13B/65B                       | q_proj,v_proj     | xverse    |
 | [Yi (1/1.5)](https://huggingface.co/01-ai)               | 6B/9B/34B                        | q_proj,v_proj     | yi        |
+| [Yi-VL](https://huggingface.co/01-ai)                    | 6B                               | q_proj,v_proj     | yi_vl     |
 | [Yuan](https://huggingface.co/IEITYuan)                  | 2B/51B/102B                      | q_proj,v_proj     | yuan      |
 
 > [!NOTE]

src/llmtuner/data/template.py

@@ -868,6 +868,7 @@ _register_template(
         "仔细阅读所有的图像，并对人类的问题做出信息丰富、有帮助、详细的和礼貌的回答。\n\n"
     ),
     stop_words=["###"],
+    efficient_eos=True,
 )
 
 

src/llmtuner/extras/constants.py

@@ -26,8 +26,6 @@ LAYERNORM_NAMES = {"norm", "ln"}
 
 METHODS = ["full", "freeze", "lora"]
 
-MLLM_LIST = ["LLaVA1.5"]
-
 MOD_SUPPORTED_MODELS = ["bloom", "falcon", "gemma", "llama", "mistral", "mixtral", "phi", "starcoder2"]
 
 PEFT_METHODS = ["lora"]
@@ -59,6 +57,8 @@ V_HEAD_WEIGHTS_NAME = "value_head.bin"
 
 V_HEAD_SAFE_WEIGHTS_NAME = "value_head.safetensors"
 
+VISION_MODELS = set()
+
 
 class DownloadSource(str, Enum):
     DEFAULT = "hf"
@@ -69,6 +69,7 @@ def register_model_group(
     models: Dict[str, Dict[DownloadSource, str]],
     module: Optional[str] = None,
     template: Optional[str] = None,
+    vision: bool = False,
 ) -> None:
     prefix = None
     for name, path in models.items():
@@ -81,6 +82,8 @@ def register_model_group(
         DEFAULT_MODULE[prefix] = module
     if template is not None:
         DEFAULT_TEMPLATE[prefix] = template
+    if vision:
+        VISION_MODELS.add(prefix)
 
 
 register_model_group(
@@ -599,6 +602,7 @@ register_model_group(
         },
     },
     template="vicuna",
+    vision=True,
 )
 
 
@@ -1206,6 +1210,17 @@ register_model_group(
 )
 
 
+register_model_group(
+    models={
+        "YiVL-6B-Chat": {
+            DownloadSource.DEFAULT: "BUAADreamer/Yi-VL-6B-hf",
+        },
+    },
+    template="yi_vl",
+    vision=True,
+)
+
+
 register_model_group(
     models={
         "Yuan2-2B-Chat": {

src/llmtuner/webui/common.py

@@ -10,11 +10,11 @@ from ..extras.constants import (
     DATA_CONFIG,
     DEFAULT_MODULE,
     DEFAULT_TEMPLATE,
-    MLLM_LIST,
     PEFT_METHODS,
     STAGES_USE_PAIR_DATA,
     SUPPORTED_MODELS,
     TRAINING_STAGES,
+    VISION_MODELS,
     DownloadSource,
 )
 from ..extras.logging import get_logger
@@ -112,7 +112,7 @@ def get_template(model_name: str) -> str:
 
 
 def get_visual(model_name: str) -> bool:
-    return get_prefix(model_name) in MLLM_LIST
+    return get_prefix(model_name) in VISION_MODELS
 
 
 def list_adapters(model_name: str, finetuning_type: str) -> "gr.Dropdown":