support switching base models

LLaMA_LoRA.ipynb

@@ -289,7 +289,8 @@
         "\n",
         "# Set Configs\n",
         "from llama_lora.llama_lora.globals import Global\n",
-        "Global.default_base_model_name = base_model\n",
+        "Global.default_base_model_name = Global.base_model_name = base_model\n",
+        "Global.base_model_choices = [base_model]\n",
         "data_dir_realpath = !realpath ./data\n",
         "Global.data_dir = data_dir_realpath[0]\n",
         "Global.load_8bit = True\n",

app.py

@@ -14,6 +14,7 @@ from llama_lora.utils.data import init_data_dir
 def main(
     base_model: str = "",
     data_dir: str = "",
+    base_model_choices: str = "",
     # Allows to listen on all interfaces by providing '0.0.0.0'.
     server_name: str = "127.0.0.1",
     share: bool = False,
@@ -29,6 +30,9 @@ def main(
 
     :param base_model: (required) The name of the default base model to use.
     :param data_dir: (required) The path to the directory to store data.
+
+    :param base_model_choices: Base model selections to display on the UI, seperated by ",". For example: 'decapoda-research/llama-7b-hf,nomic-ai/gpt4all-j'.
+
     :param server_name: Allows to listen on all interfaces by providing '0.0.0.0'.
     :param share: Create a public Gradio URL.
 
@@ -46,7 +50,16 @@ def main(
         data_dir
     ), "Please specify a --data_dir, e.g. --data_dir='./data'"
 
-    Global.default_base_model_name = base_model
+    Global.default_base_model_name = Global.base_model_name = base_model
+
+    if base_model_choices:
+        base_model_choices = base_model_choices.split(',')
+        base_model_choices = [name.strip() for name in base_model_choices]
+        Global.base_model_choices = base_model_choices
+
+    if base_model not in Global.base_model_choices:
+        Global.base_model_choices = [base_model] + Global.base_model_choices
+
     Global.data_dir = os.path.abspath(data_dir)
     Global.load_8bit = load_8bit
 

download_base_model.py

@@ -0,0 +1,32 @@
+import fire
+
+from llama_lora.models import get_new_base_model, clear_cache
+
+
+def main(
+    base_model_names: str = "",
+):
+    '''
+    Download and cache base models form Hugging Face.
+
+    :param base_model_names: Names of the base model you want to download, seperated by ",". For example: 'decapoda-research/llama-7b-hf,nomic-ai/gpt4all-j'.
+    '''
+
+    assert (
+        base_model_names
+    ), "Please specify --base_model_names, e.g. --base_model_names='decapoda-research/llama-7b-hf,nomic-ai/gpt4all-j'"
+
+    base_model_names = base_model_names.split(',')
+    base_model_names = [name.strip() for name in base_model_names]
+
+    print(f"Base models: {', '.join(base_model_names)}.")
+
+    for name in base_model_names:
+        print(f"Preparing {name}...")
+        get_new_base_model(name)
+        clear_cache()
+
+    print("Done.")
+
+if __name__ == "__main__":
+    fire.Fire(main)

llama_lora/globals.py

@@ -17,6 +17,8 @@ class Global:
     load_8bit: bool = False
 
     default_base_model_name: str = ""
+    base_model_name: str = ""
+    base_model_choices: List[str] = []
 
     # Functions
     train_fn: Any = train

llama_lora/models.py

@@ -2,9 +2,10 @@ import os
 import sys
 import gc
 import json
+import re
 
 import torch
-from transformers import LlamaForCausalLM, LlamaTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaTokenizer
 from peft import PeftModel
 
 from .globals import Global
@@ -29,7 +30,7 @@ def get_new_base_model(base_model_name):
     device = get_device()
 
     if device == "cuda":
-        model = LlamaForCausalLM.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             base_model_name,
             load_in_8bit=Global.load_8bit,
             torch_dtype=torch.float16,
@@ -38,20 +39,22 @@ def get_new_base_model(base_model_name):
             device_map={'': 0},
         )
     elif device == "mps":
-        model = LlamaForCausalLM.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             base_model_name,
             device_map={"": device},
             torch_dtype=torch.float16,
         )
     else:
-        model = LlamaForCausalLM.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             base_model_name, device_map={"": device}, low_cpu_mem_usage=True
         )
 
     tokenizer = get_tokenizer(base_model_name)
-    model.config.pad_token_id = tokenizer.pad_token_id = 0
-    model.config.bos_token_id = tokenizer.bos_token_id = 1
-    model.config.eos_token_id = tokenizer.eos_token_id = 2
+
+    if re.match("[^/]+/llama", base_model_name):
+        model.config.pad_token_id = tokenizer.pad_token_id = 0
+        model.config.bos_token_id = tokenizer.bos_token_id = 1
+        model.config.eos_token_id = tokenizer.eos_token_id = 2
 
     return model
 
@@ -64,7 +67,14 @@ def get_tokenizer(base_model_name):
     if loaded_tokenizer:
         return loaded_tokenizer
 
-    tokenizer = LlamaTokenizer.from_pretrained(base_model_name)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+    except Exception as e:
+        if 'LLaMATokenizer' in str(e):
+            tokenizer = LlamaTokenizer.from_pretrained(base_model_name)
+        else:
+            raise e
+
     Global.loaded_tokenizers.set(base_model_name, tokenizer)
 
     return tokenizer
@@ -137,9 +147,10 @@ def get_model(
                 device_map={"": device},
             )
 
-    model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0
-    model.config.bos_token_id = 1
-    model.config.eos_token_id = 2
+    if re.match("[^/]+/llama", base_model_name):
+        model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0
+        model.config.bos_token_id = 1
+        model.config.eos_token_id = 2
 
     if not Global.load_8bit:
         model.half()  # seems to fix bugs for some users.

llama_lora/ui/finetune_ui.py

@@ -299,7 +299,7 @@ def do_train(
     progress=gr.Progress(track_tqdm=should_training_progress_track_tqdm),
 ):
     try:
-        base_model_name = Global.default_base_model_name
+        base_model_name = Global.base_model_name
         output_dir = os.path.join(Global.data_dir, "lora_models", model_name)
         if os.path.exists(output_dir):
             if (not os.path.isdir(output_dir)) or os.path.exists(os.path.join(output_dir, 'adapter_config.json')):

llama_lora/ui/inference_ui.py

@@ -22,7 +22,7 @@ inference_output_lines = 12
 
 
 def prepare_inference(lora_model_name, progress=gr.Progress(track_tqdm=True)):
-    base_model_name = Global.default_base_model_name
+    base_model_name = Global.base_model_name
 
     try:
         get_tokenizer(base_model_name)
@@ -48,7 +48,7 @@ def do_inference(
     show_raw=False,
     progress=gr.Progress(track_tqdm=True),
 ):
-    base_model_name = Global.default_base_model_name
+    base_model_name = Global.base_model_name
 
     try:
         if Global.generation_force_stopped_at is not None:
@@ -257,7 +257,7 @@ def reload_selections(current_lora_model, current_prompt_template):
     current_prompt_template = current_prompt_template or next(
         iter(available_template_names_with_none), None)
 
-    default_lora_models = ["tloen/alpaca-lora-7b"]
+    default_lora_models = []
     available_lora_models = default_lora_models + get_available_lora_model_names()
     available_lora_models = available_lora_models + ["None"]
 
@@ -283,8 +283,12 @@ def handle_prompt_template_change(prompt_template, lora_model):
         "", visible=False)
     lora_mode_info = get_info_of_available_lora_model(lora_model)
     if lora_mode_info and isinstance(lora_mode_info, dict):
+        model_base_model = lora_mode_info.get("base_model")
         model_prompt_template = lora_mode_info.get("prompt_template")
-        if model_prompt_template and model_prompt_template != prompt_template:
+        if model_base_model and model_base_model != Global.base_model_name:
+            model_prompt_template_message_update = gr.Markdown.update(
+                f"⚠️ This model was trained on top of base model `{model_base_model}`, it might not work properly with the selected base model `{Global.base_model_name}`.", visible=True)
+        elif model_prompt_template and model_prompt_template != prompt_template:
             model_prompt_template_message_update = gr.Markdown.update(
                 f"This model was trained with prompt template `{model_prompt_template}`.", visible=True)
 
@@ -331,7 +335,7 @@ def inference_ui():
                 lora_model = gr.Dropdown(
                     label="LoRA Model",
                     elem_id="inference_lora_model",
-                    value="tloen/alpaca-lora-7b",
+                    value="None",
                     allow_custom_value=True,
                 )
             prompt_template = gr.Dropdown(
@@ -461,6 +465,8 @@ def inference_ui():
                             interactive=False,
                             elem_id="inference_raw_output")
 
+        reload_selected_models_btn = gr.Button("", elem_id="inference_reload_selected_models_btn")
+
         show_raw_change_event = show_raw.change(
             fn=lambda show_raw: gr.Accordion.update(visible=show_raw),
             inputs=[show_raw],
@@ -482,6 +488,14 @@ def inference_ui():
                 variable_0, variable_1, variable_2, variable_3, variable_4, variable_5, variable_6, variable_7])
         things_that_might_timeout.append(prompt_template_change_event)
 
+        reload_selected_models_btn_event = reload_selected_models_btn.click(
+            fn=handle_prompt_template_change,
+            inputs=[prompt_template, lora_model],
+            outputs=[
+                model_prompt_template_message,
+                variable_0, variable_1, variable_2, variable_3, variable_4, variable_5, variable_6, variable_7])
+        things_that_might_timeout.append(reload_selected_models_btn_event)
+
         lora_model_change_event = lora_model.change(
             fn=handle_lora_model_change,
             inputs=[lora_model, prompt_template],
@@ -538,7 +552,7 @@ def inference_ui():
 
         // Workaround default value not shown.
         document.querySelector('#inference_lora_model input').value =
-          'tloen/alpaca-lora-7b';
+          'None';
       }, 100);
 
       // Add tooltips
@@ -682,6 +696,30 @@ def inference_ui():
         }, 500);
       }, 0);
 
+      // Reload model selection on possible base model change.
+      setTimeout(function () {
+        const elem = document.getElementById('main_page_tabs_container');
+        if (!elem) return;
+
+        let prevClassList = [];
+
+        new MutationObserver(function (mutationsList, observer) {
+          const currentPrevClassList = prevClassList;
+          const currentClassList = Array.from(elem.classList);
+          prevClassList = Array.from(elem.classList);
+
+          if (!currentPrevClassList.includes('hide')) return;
+          if (currentClassList.includes('hide')) return;
+
+          const inference_reload_selected_models_btn_elem = document.getElementById('inference_reload_selected_models_btn');
+
+          if (inference_reload_selected_models_btn_elem) inference_reload_selected_models_btn_elem.click();
+        }).observe(elem, {
+          attributes: true,
+          attributeFilter: ['class'],
+        });
+      }, 0);
+
       // Debounced updating the prompt preview.
       setTimeout(function () {
         function debounce(func, wait) {

llama_lora/ui/main_page.py

@@ -17,25 +17,50 @@ def main_page():
             css=main_page_custom_css(),
     ) as main_page_blocks:
         with gr.Column(elem_id="main_page_content"):
-            gr.Markdown(f"""
-                <h1 class="app_title_text">{title}</h1> <wbr />
-                <h2 class="app_subtitle_text">{Global.ui_subtitle}</h2>
-                """)
-            with gr.Tab("Inference"):
-                inference_ui()
-            with gr.Tab("Fine-tuning"):
-                finetune_ui()
-            with gr.Tab("Tokenizer"):
-                tokenizer_ui()
-            info = []
-            if Global.version:
-                info.append(f"LLaMA-LoRA Tuner `{Global.version}`")
-            info.append(f"Base model: `{Global.default_base_model_name}`")
-            if Global.ui_show_sys_info:
-                info.append(f"Data dir: `{Global.data_dir}`")
-            gr.Markdown(f"""
-                <small>{"&nbsp;&nbsp;·&nbsp;&nbsp;".join(info)}</small>
-                """)
+            with gr.Row():
+                gr.Markdown(
+                    f"""
+                    <h1 class="app_title_text">{title}</h1> <wbr />
+                    <h2 class="app_subtitle_text">{Global.ui_subtitle}</h2>
+                    """,
+                    elem_id="page_title",
+                )
+                global_base_model_select = gr.Dropdown(
+                    label="Base Model",
+                    elem_id="global_base_model_select",
+                    choices=Global.base_model_choices,
+                    value=lambda: Global.base_model_name,
+                    allow_custom_value=True,
+                )
+            # global_base_model_select_loading_status = gr.Markdown("", elem_id="global_base_model_select_loading_status")
+
+            with gr.Column(elem_id="main_page_tabs_container") as main_page_tabs_container:
+                with gr.Tab("Inference"):
+                    inference_ui()
+                with gr.Tab("Fine-tuning"):
+                    finetune_ui()
+                with gr.Tab("Tokenizer"):
+                    tokenizer_ui()
+            please_select_a_base_model_message = gr.Markdown("Please select a base model.", visible=False)
+            current_base_model_hint = gr.Markdown(lambda: Global.base_model_name, elem_id="current_base_model_hint")
+            foot_info = gr.Markdown(get_foot_info)
+
+    global_base_model_select.change(
+        fn=pre_handle_change_base_model,
+        inputs=[],
+        outputs=[main_page_tabs_container]
+    ).then(
+        fn=handle_change_base_model,
+        inputs=[global_base_model_select],
+        outputs=[
+            main_page_tabs_container,
+            please_select_a_base_model_message,
+            current_base_model_hint,
+            # global_base_model_select_loading_status,
+            foot_info
+        ]
+    )
+
     main_page_blocks.load(_js=f"""
     function () {{
         {popperjs_core_code()}
@@ -61,6 +86,17 @@ def main_page():
           });
           handle_gradio_container_element_class_change();
         }, 500);
+    """ + """
+        setTimeout(function () {
+          // Workaround default value not shown.
+          const current_base_model_hint_elem = document.querySelector('#current_base_model_hint > p');
+          if (!current_base_model_hint_elem) return;
+
+          const base_model_name = current_base_model_hint_elem.innerText;
+          document.querySelector('#global_base_model_select input').value = base_model_name;
+          document.querySelector('#global_base_model_select').classList.add('show');
+        }, 3200);
+    """ + """
     }
     """)
 
@@ -127,12 +163,77 @@ def main_page_custom_css():
         display: none;
     }
 
+    #page_title {
+        flex-grow: 3;
+    }
+    #global_base_model_select {
+        position: relative;
+        align-self: center;
+        min-width: 250px;
+        padding: 2px 2px;
+        border: 0;
+        box-shadow: none;
+        opacity: 0;
+        pointer-events: none;
+    }
+    #global_base_model_select.show {
+        opacity: 1;
+        pointer-events: auto;
+    }
+    #global_base_model_select label .wrap-inner {
+        padding: 2px 8px;
+    }
+    #global_base_model_select label span {
+        margin-bottom: 2px;
+        font-size: 80%;
+        position: absolute;
+        top: -14px;
+        left: 8px;
+        opacity: 0;
+    }
+    #global_base_model_select:hover label span {
+        opacity: 1;
+    }
+
+    #global_base_model_select_loading_status {
+        position: absolute;
+        pointer-events: none;
+        top: 0;
+        left: 0;
+        right: 0;
+        bottom: 0;
+    }
+    #global_base_model_select_loading_status > .wrap:not(.hide) {
+        z-index: 9999;
+        position: absolute;
+        top: 112px !important;
+        bottom: 0 !important;
+        max-height: none;
+        background: var(--background-fill-primary);
+        opacity: 0.8;
+    }
+
+    #current_base_model_hint  {
+        display: none;
+    }
+
     #main_page_content > .tabs > .tab-nav * {
         font-size: 1rem;
         font-weight: 700;
         /* text-transform: uppercase; */
     }
 
+    #inference_reload_selected_models_btn {
+        position: absolute;
+        top: 0;
+        left: 0;
+        width: 0;
+        height: 0;
+        padding: 0;
+        opacity: 0;
+        pointer-events: none;
+    }
+
     #inference_lora_model_group {
         border-radius: var(--block-radius);
         background: var(--block-background-fill);
@@ -147,7 +248,8 @@ def main_page_custom_css():
         position: absolute;
         bottom: 8px;
         left: 20px;
-        z-index: 1;
+        z-index: 61;
+        width: 999px;
         font-size: 12px;
         opacity: 0.7;
     }
@@ -515,3 +617,28 @@ def main_page_custom_css():
     .tippy-box[data-animation=scale-subtle][data-placement^=top]{transform-origin:bottom}.tippy-box[data-animation=scale-subtle][data-placement^=bottom]{transform-origin:top}.tippy-box[data-animation=scale-subtle][data-placement^=left]{transform-origin:right}.tippy-box[data-animation=scale-subtle][data-placement^=right]{transform-origin:left}.tippy-box[data-animation=scale-subtle][data-state=hidden]{transform:scale(.8);opacity:0}
     """
     return css
+
+
+def pre_handle_change_base_model():
+    return gr.Column.update(visible=False)
+
+
+def handle_change_base_model(selected_base_model_name):
+    Global.base_model_name = selected_base_model_name
+
+    if Global.base_model_name:
+        return gr.Column.update(visible=True), gr.Markdown.update(visible=False), Global.base_model_name, get_foot_info()
+
+    return gr.Column.update(visible=False), gr.Markdown.update(visible=True), Global.base_model_name, get_foot_info()
+
+
+def get_foot_info():
+    info = []
+    if Global.version:
+        info.append(f"LLaMA-LoRA Tuner `{Global.version}`")
+    info.append(f"Base model: `{Global.base_model_name}`")
+    if Global.ui_show_sys_info:
+        info.append(f"Data dir: `{Global.data_dir}`")
+    return f"""\
+        <small>{"&nbsp;&nbsp;·&nbsp;&nbsp;".join(info)}</small>
+        """

lora_models/alpaca-lora-7b/info.json

@@ -0,0 +1,6 @@
+{
+  "hf_model_name": "tloen/alpaca-lora-7b",
+  "load_from_hf": true,
+  "base_model": "decapoda-research/llama-7b-hf",
+  "prompt_template": "alpaca"
+}