support switching to custom tokenizer

llama_lora/globals.py

@@ -18,6 +18,7 @@ class Global:
 
     default_base_model_name: str = ""
     base_model_name: str = ""
+    tokenizer_name = None
     base_model_choices: List[str] = []
 
     trust_remote_code = False

llama_lora/lib/csv_logger.py

@@ -25,8 +25,8 @@ class CSVLogger(FlaggingCallback):
 
     def setup(
         self,
-        components: List[Any],
-        flagging_dir: str | Path,
+        components,
+        flagging_dir,
     ):
         self.components = components
         self.flagging_dir = flagging_dir
@@ -36,7 +36,7 @@ class CSVLogger(FlaggingCallback):
         self,
         flag_data: List[Any],
         flag_option: str = "",
-        username: str | None = None,
+        username=None,
         filename="log.csv",
     ) -> int:
         flagging_dir = self.flagging_dir

llama_lora/ui/finetune_ui.py

@@ -306,6 +306,7 @@ def do_train(
 ):
     try:
         base_model_name = Global.base_model_name
+        tokenizer_name = Global.tokenizer_name or Global.base_model_name
 
         resume_from_checkpoint = None
         if continue_from_model == "-" or continue_from_model == "None":
@@ -445,7 +446,7 @@ Train data (first 10):
         Global.should_stop_training = False
 
         base_model = get_new_base_model(base_model_name)
-        tokenizer = get_tokenizer(base_model_name)
+        tokenizer = get_tokenizer(tokenizer_name)
 
         # Do not let other tqdm iterations interfere the progress reporting after training starts.
         # progress.track_tqdm = False  # setting this dynamically is not working, determining if track_tqdm should be enabled based on GPU cores at start instead.

llama_lora/ui/inference_ui.py

@@ -33,9 +33,10 @@ class LoggingItem:
 
 def prepare_inference(lora_model_name, progress=gr.Progress(track_tqdm=True)):
     base_model_name = Global.base_model_name
+    tokenizer_name = Global.tokenizer_name or Global.base_model_name
 
     try:
-        get_tokenizer(base_model_name)
+        get_tokenizer(tokenizer_name)
         get_model(base_model_name, lora_model_name)
         return ("", "", gr.Textbox.update(visible=False))
 

llama_lora/ui/main_page.py

@@ -25,13 +25,29 @@ def main_page():
                     """,
                     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,
-                )
+                with gr.Column(elem_id="global_base_model_select_group"):
+                    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,
+                    )
+                    use_custom_tokenizer_btn = gr.Button(
+                        "Use custom tokenizer",
+                        elem_id="use_custom_tokenizer_btn")
+                    global_tokenizer_select = gr.Dropdown(
+                        label="Tokenizer",
+                        elem_id="global_tokenizer_select",
+                        # choices=[],
+                        value=lambda: Global.base_model_name,
+                        visible=False,
+                        allow_custom_value=True,
+                    )
+                    use_custom_tokenizer_btn.click(
+                        fn=lambda: gr.Dropdown.update(visible=True),
+                        inputs=None,
+                        outputs=[global_tokenizer_select])
             # 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:
@@ -41,13 +57,17 @@ def main_page():
                     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")
+            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")
+            current_tokenizer_hint = gr.Markdown(
+                lambda: Global.tokenizer_name, elem_id="current_tokenizer_hint")
             foot_info = gr.Markdown(get_foot_info)
 
     global_base_model_select.change(
         fn=pre_handle_change_base_model,
-        inputs=[],
+        inputs=[global_base_model_select],
         outputs=[main_page_tabs_container]
     ).then(
         fn=handle_change_base_model,
@@ -56,11 +76,27 @@ def main_page():
             main_page_tabs_container,
             please_select_a_base_model_message,
             current_base_model_hint,
+            current_tokenizer_hint,
             # global_base_model_select_loading_status,
             foot_info
         ]
     )
 
+    global_tokenizer_select.change(
+        fn=pre_handle_change_tokenizer,
+        inputs=[global_tokenizer_select],
+        outputs=[main_page_tabs_container]
+    ).then(
+        fn=handle_change_tokenizer,
+        inputs=[global_tokenizer_select],
+        outputs=[
+            global_tokenizer_select,
+            main_page_tabs_container,
+            current_tokenizer_hint,
+            foot_info
+        ]
+    )
+
     main_page_blocks.load(_js=f"""
     function () {{
         {popperjs_core_code()}
@@ -95,6 +131,15 @@ def main_page():
           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');
+
+          const current_tokenizer_hint_elem = document.querySelector('#current_tokenizer_hint > p');
+          const tokenizer_name = current_tokenizer_hint_elem && current_tokenizer_hint_elem.innerText;
+
+          if (tokenizer_name && tokenizer_name !== base_model_name) {
+            document.querySelector('#global_tokenizer_select input').value = tokenizer_name;
+            const btn = document.getElementById('use_custom_tokenizer_btn');
+            if (btn) btn.click();
+          }
         }, 3200);
     """ + """
     }
@@ -209,13 +254,21 @@ def main_page_custom_css():
     #page_title {
         flex-grow: 3;
     }
-    #global_base_model_select {
+    #global_base_model_select_group,
+    #global_base_model_select,
+    #global_tokenizer_select {
         position: relative;
         align-self: center;
-        min-width: 250px;
+        min-width: 250px !important;
+    }
+    #global_base_model_select,
+    #global_tokenizer_select {
+        position: relative;
         padding: 2px 2px;
         border: 0;
         box-shadow: none;
+    }
+    #global_base_model_select {
         opacity: 0;
         pointer-events: none;
     }
@@ -223,10 +276,12 @@ def main_page_custom_css():
         opacity: 1;
         pointer-events: auto;
     }
-    #global_base_model_select label .wrap-inner {
+    #global_base_model_select label .wrap-inner,
+    #global_tokenizer_select label .wrap-inner {
         padding: 2px 8px;
     }
-    #global_base_model_select label span {
+    #global_base_model_select label span,
+    #global_tokenizer_select label span {
         margin-bottom: 2px;
         font-size: 80%;
         position: absolute;
@@ -234,9 +289,28 @@ def main_page_custom_css():
         left: 8px;
         opacity: 0;
     }
-    #global_base_model_select:hover label span {
+    #global_base_model_select_group:hover label span,
+    #global_base_model_select:hover label span,
+    #global_tokenizer_select:hover label span {
         opacity: 1;
     }
+    #use_custom_tokenizer_btn {
+        position: absolute;
+        top: -16px;
+        right: 10px;
+        border: 0 !important;
+        width: auto !important;
+        background: transparent !important;
+        box-shadow: none !important;
+        padding: 0 !important;
+        font-weight: 100 !important;
+        text-decoration: underline;
+        font-size: 12px !important;
+        opacity: 0;
+    }
+    #global_base_model_select_group:hover #use_custom_tokenizer_btn {
+        opacity: 0.3;
+    }
 
     #global_base_model_select_loading_status {
         position: absolute;
@@ -260,7 +334,7 @@ def main_page_custom_css():
         background: var(--block-background-fill);
     }
 
-    #current_base_model_hint  {
+    #current_base_model_hint, #current_tokenizer_hint {
         display: none;
     }
 
@@ -754,24 +828,61 @@ def main_page_custom_css():
     return css
 
 
-def pre_handle_change_base_model():
-    return gr.Column.update(visible=False)
+def pre_handle_change_base_model(selected_base_model_name):
+    if Global.base_model_name != selected_base_model_name:
+        return gr.Column.update(visible=False)
+    if Global.tokenizer_name and Global.tokenizer_name != selected_base_model_name:
+        return gr.Column.update(visible=False)
+    return gr.Column.update(visible=True)
 
 
 def handle_change_base_model(selected_base_model_name):
     Global.base_model_name = selected_base_model_name
+    Global.tokenizer_name = selected_base_model_name
 
+    is_base_model_selected = False
     if Global.base_model_name:
-        return gr.Column.update(visible=True), gr.Markdown.update(visible=False), Global.base_model_name, get_foot_info()
+        is_base_model_selected = True
+
+    return (
+        gr.Column.update(visible=is_base_model_selected),
+        gr.Markdown.update(visible=not is_base_model_selected),
+        Global.base_model_name,
+        Global.tokenizer_name,
+        get_foot_info())
+
+
+def pre_handle_change_tokenizer(selected_tokenizer_name):
+    if Global.tokenizer_name != selected_tokenizer_name:
+        return gr.Column.update(visible=False)
+    return gr.Column.update(visible=True)
+
 
-    return gr.Column.update(visible=False), gr.Markdown.update(visible=True), Global.base_model_name, get_foot_info()
+def handle_change_tokenizer(selected_tokenizer_name):
+    Global.tokenizer_name = selected_tokenizer_name
+
+    show_tokenizer_select = True
+    if not Global.tokenizer_name:
+        show_tokenizer_select = False
+    if Global.tokenizer_name == Global.base_model_name:
+        show_tokenizer_select = False
+
+    return (
+        gr.Dropdown.update(visible=show_tokenizer_select),
+        gr.Column.update(visible=True),
+        Global.tokenizer_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.base_model_name:
+        info.append(f"Base model: `{Global.base_model_name}`")
+    if Global.tokenizer_name and Global.tokenizer_name != Global.base_model_name:
+        info.append(f"Tokenizer: `{Global.tokenizer_name}`")
     if Global.ui_show_sys_info:
         info.append(f"Data dir: `{Global.data_dir}`")
     return f"""\

llama_lora/ui/tokenizer_ui.py

@@ -7,12 +7,14 @@ from ..models import get_tokenizer
 
 
 def handle_decode(encoded_tokens_json):
-    base_model_name = Global.base_model_name
+    # base_model_name = Global.base_model_name
+    tokenizer_name = Global.tokenizer_name or Global.base_model_name
+
     try:
         encoded_tokens = json.loads(encoded_tokens_json)
         if Global.ui_dev_mode:
             return f"Not actually decoding tokens in UI dev mode.", gr.Markdown.update("", visible=False)
-        tokenizer = get_tokenizer(base_model_name)
+        tokenizer = get_tokenizer(tokenizer_name)
         decoded_tokens = tokenizer.decode(encoded_tokens)
         return decoded_tokens, gr.Markdown.update("", visible=False)
     except Exception as e:
@@ -20,11 +22,13 @@ def handle_decode(encoded_tokens_json):
 
 
 def handle_encode(decoded_tokens):
-    base_model_name = Global.base_model_name
+    # base_model_name = Global.base_model_name
+    tokenizer_name = Global.tokenizer_name or Global.base_model_name
+
     try:
         if Global.ui_dev_mode:
             return f"[\"Not actually encoding tokens in UI dev mode.\"]", gr.Markdown.update("", visible=False)
-        tokenizer = get_tokenizer(base_model_name)
+        tokenizer = get_tokenizer(tokenizer_name)
         result = tokenizer(decoded_tokens)
         encoded_tokens_json = json.dumps(result['input_ids'], indent=2)
         return encoded_tokens_json, gr.Markdown.update("", visible=False)