Merge branch 'main' into hf-ui-demo

.gitignore

@@ -3,4 +3,5 @@ __pycache__/
 /venv
 .vscode
 
+/wandb
 /data

LLaMA_LoRA.ipynb

@@ -60,12 +60,20 @@
         "# @title A small workaround { display-mode: \"form\" }\n",
         "# @markdown Don't panic if you see an error here. Just click the `RESTART RUNTIME` button in the output below, then Run All again.\n",
         "# @markdown The error will disappear on the next run.\n",
-        "!pip install Pillow==9.3.0\n",
+        "!pip install Pillow==9.3.0 numpy==1.23.5\n",
+        "\n",
         "import PIL\n",
         "major, minor = map(float, PIL.__version__.split(\".\")[:2])\n",
         "version_float = major + minor / 10**len(str(minor))\n",
-        "print(version_float)\n",
+        "print('PIL', version_float)\n",
         "if version_float < 9.003:\n",
+        "    raise Exception(\"Restart the runtime by clicking the 'RESTART RUNTIME' button above (or Runtime > Restart Runtime).\")\n",
+        "\n",
+        "import numpy\n",
+        "major, minor = map(float, numpy.__version__.split(\".\")[:2])\n",
+        "version_float = major + minor / 10**len(str(minor))\n",
+        "print('numpy', version_float)\n",
+        "if version_float < 1.0023:\n",
         "    raise Exception(\"Restart the runtime by clicking the 'RESTART RUNTIME' button above (or Runtime > Restart Runtime).\")"
       ],
       "metadata": {
@@ -281,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",

README.md

@@ -34,8 +34,8 @@ Making evaluating and fine-tuning LLaMA models with low-rank adaptation (LoRA) e
 
 * **[1-click up and running in Google Colab](#run-on-google-colab)** with a standard GPU runtime.
   * Loads and stores data in Google Drive.
-* Evaluate various LLaMA LoRA models stored in your folder or from Hugging Face.<br /><a href="https://youtu.be/A3kb4VkDWyY"><img width="640px" src="https://user-images.githubusercontent.com/3784687/230272844-09f7a35b-46bf-4101-b15d-4ddf243b8bef.gif" /></a>
-* Fine-tune LLaMA models with different prompt templates and training dataset format.<br /><a href="https://youtu.be/5Db9U8PsaUk"><img width="640px" src="https://user-images.githubusercontent.com/3784687/230277315-9a91d983-1690-4594-9d54-912eda8963ee.gif" /></a>
+* Evaluate various LLaMA LoRA models stored in your folder or from Hugging Face.<br /><a href="https://youtu.be/IoEMgouZ5xU"><img width="640px" src="https://user-images.githubusercontent.com/3784687/231023326-f28c84e2-df74-4179-b0ac-c25c4e8ca001.gif" /></a>
+* Fine-tune LLaMA models with different prompt templates and training dataset format.<br /><a href="https://youtu.be/IoEMgouZ5xU?t=60"><img width="640px" src="https://user-images.githubusercontent.com/3784687/231026640-b5cf5c79-9fe9-430b-8d4e-7346eb9567ad.gif" /></a>
   * Load JSON and JSONL datasets from your folder, or even paste plain text directly into the UI.
   * Supports Stanford Alpaca [seed_tasks](https://github.com/tatsu-lab/stanford_alpaca/blob/main/seed_tasks.jsonl), [alpaca_data](https://github.com/tatsu-lab/stanford_alpaca/blob/main/alpaca_data.json) and [OpenAI "prompt"-"completion"](https://platform.openai.com/docs/guides/fine-tuning/data-formatting) format.
   * Use prompt templates to keep your dataset DRY.
@@ -51,6 +51,8 @@ There are various ways to run this app:
 
 ### Run On Google Colab
 
+*See [video](https://youtu.be/lByYOMdy9h4) for step-by-step instructions.*
+
 Open [this Colab Notebook](https://colab.research.google.com/github/zetavg/LLaMA-LoRA-Tuner/blob/main/LLaMA_LoRA.ipynb) and select **Runtime > Run All** (`⌘/Ctrl+F9`).
 
 You will be prompted to authorize Google Drive access, as Google Drive will be used to store your data. See the "Config"/"Google Drive" section for settings and more info.
@@ -81,13 +83,14 @@ file_mounts:
 setup: |
   git clone https://github.com/zetavg/LLaMA-LoRA-Tuner.git llama_lora_tuner
   cd llama_lora_tuner && pip install -r requirements.lock.txt
+  pip install wandb
   cd ..
   echo 'Dependencies installed.'
 
 # Start the app.
 run: |
   echo 'Starting...'
-  python llama_lora_tuner/app.py --data_dir='/data' --base_model='decapoda-research/llama-7b-hf' --share
+  python llama_lora_tuner/app.py --data_dir='/data' --wandb_api_key "$([ -f /data/secrets/wandb_api_key ] && cat /data/secrets/wandb_api_key | tr -d '\n')" --base_model='decapoda-research/llama-7b-hf' --share
 ```
 
 Then launch a cluster to run the task:
@@ -135,6 +138,11 @@ For more options, see `python app.py --help`.
 </details>
 
 
+## Usage
+
+See [video on YouTube](https://youtu.be/IoEMgouZ5xU).
+
+
 ## Acknowledgements
 
 * https://github.com/tloen/alpaca-lora

app.py

@@ -5,21 +5,41 @@ import fire
 import gradio as gr
 
 from llama_lora.globals import Global
+from llama_lora.models import prepare_base_model
 from llama_lora.ui.main_page import main_page, get_page_title, main_page_custom_css
 from llama_lora.utils.data import init_data_dir
 
 
+
 def main(
-    load_8bit: bool = False,
     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,
     skip_loading_base_model: bool = False,
+    load_8bit: bool = False,
     ui_show_sys_info: bool = True,
     ui_dev_mode: bool = False,
+    wandb_api_key: str = "",
+    wandb_project: str = "",
 ):
+    '''
+    Start the LLaMA-LoRA Tuner UI.
+
+    :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.
+
+    :param wandb_api_key: The API key for Weights & Biases. Setting either this or `wandb_project` will enable Weights & Biases.
+    :param wandb_project: The default project name for Weights & Biases. Setting either this or `wandb_api_key` will enable Weights & Biases.
+    '''
+
     base_model = base_model or os.environ.get("LLAMA_LORA_BASE_MODEL", "")
     data_dir = data_dir or os.environ.get("LLAMA_LORA_DATA_DIR", "")
     assert (
@@ -30,16 +50,35 @@ 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
 
+    if len(wandb_api_key) > 0:
+        Global.enable_wandb = True
+        Global.wandb_api_key = wandb_api_key
+    if len(wandb_project) > 0:
+        Global.enable_wandb = True
+        Global.wandb_project = wandb_project
+
     Global.ui_dev_mode = ui_dev_mode
     Global.ui_show_sys_info = ui_show_sys_info
 
     os.makedirs(data_dir, exist_ok=True)
     init_data_dir()
 
+    if (not skip_loading_base_model) and (not ui_dev_mode):
+        prepare_base_model(base_model)
+
     with gr.Blocks(title=get_page_title(), css=main_page_custom_css()) as demo:
         main_page()
 

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
@@ -40,6 +42,11 @@ class Global:
     gpu_total_cores = None  # GPU total cores
     gpu_total_memory = None
 
+    # WandB
+    enable_wandb = False
+    wandb_api_key = None
+    default_wandb_project = "llama-lora-tuner"
+
     # UI related
     ui_title: str = "LLaMA-LoRA Tuner"
     ui_emoji: str = "🦙🎛️"

llama_lora/lib/finetune.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import importlib
 from typing import Any, List
 
 import json
@@ -32,7 +33,7 @@ def train(
     num_train_epochs: int = 3,
     learning_rate: float = 3e-4,
     cutoff_len: int = 256,
-    val_set_size: int = 2000,  # TODO: use percentage
+    val_set_size: int = 2000,
     # lora hyperparams
     lora_r: int = 8,
     lora_alpha: int = 16,
@@ -45,13 +46,78 @@ def train(
     train_on_inputs: bool = True,  # if False, masks out inputs in loss
     group_by_length: bool = False,  # faster, but produces an odd training loss curve
     # either training checkpoint or final adapter
-    resume_from_checkpoint: str = None,
+    resume_from_checkpoint = None,
     save_steps: int = 200,
     save_total_limit: int = 3,
     logging_steps: int = 10,
     # logging
-    callbacks: List[Any] = []
+    callbacks: List[Any] = [],
+    # wandb params
+    wandb_api_key = None,
+    wandb_project: str = "",
+    wandb_group = None,
+    wandb_run_name: str = "",
+    wandb_tags: List[str] = [],
+    wandb_watch: str = "false",  # options: false | gradients | all
+    wandb_log_model: str = "true",  # options: false | true
 ):
+    # for logging
+    finetune_args = {
+        'micro_batch_size': micro_batch_size,
+        'gradient_accumulation_steps': gradient_accumulation_steps,
+        'num_train_epochs': num_train_epochs,
+        'learning_rate': learning_rate,
+        'cutoff_len': cutoff_len,
+        'val_set_size': val_set_size,
+        'lora_r': lora_r,
+        'lora_alpha': lora_alpha,
+        'lora_dropout': lora_dropout,
+        'lora_target_modules': lora_target_modules,
+        'train_on_inputs': train_on_inputs,
+        'group_by_length': group_by_length,
+        'save_steps': save_steps,
+        'save_total_limit': save_total_limit,
+        'logging_steps': logging_steps,
+    }
+    if val_set_size and val_set_size > 0:
+        finetune_args['val_set_size'] = val_set_size
+    if resume_from_checkpoint:
+        finetune_args['resume_from_checkpoint'] = resume_from_checkpoint
+
+    wandb = None
+    if wandb_api_key:
+        os.environ["WANDB_API_KEY"] = wandb_api_key
+
+    # wandb: WARNING Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to https://wandb.me/wandb-init.
+    # if wandb_project:
+    #     os.environ["WANDB_PROJECT"] = wandb_project
+    # if wandb_run_name:
+    #     os.environ["WANDB_RUN_NAME"] = wandb_run_name
+
+    if wandb_watch:
+        os.environ["WANDB_WATCH"] = wandb_watch
+    if wandb_log_model:
+        os.environ["WANDB_LOG_MODEL"] = wandb_log_model
+    use_wandb = (wandb_project and len(wandb_project) > 0) or (
+            "WANDB_PROJECT" in os.environ and len(os.environ["WANDB_PROJECT"]) > 0
+        )
+    if use_wandb:
+        os.environ['WANDB_MODE'] = "online"
+        wandb = importlib.import_module("wandb")
+        wandb.init(
+            project=wandb_project,
+            resume="auto",
+            group=wandb_group,
+            name=wandb_run_name,
+            tags=wandb_tags,
+            reinit=True,
+            magic=True,
+            config={'finetune_args': finetune_args},
+            # id=None  # used for resuming
+            )
+    else:
+        os.environ['WANDB_MODE'] = "disabled"
+
     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')):
             raise ValueError(
@@ -138,6 +204,8 @@ def train(
 
     # If train_dataset_data is a list, convert it to datasets.Dataset
     if isinstance(train_dataset_data, list):
+        with open(os.path.join(output_dir, "train_data_samples.json"), 'w') as file:
+            json.dump(list(train_dataset_data[:100]), file, indent=2)
         train_dataset_data = Dataset.from_list(train_dataset_data)
 
     if resume_from_checkpoint:
@@ -158,7 +226,7 @@ def train(
             adapters_weights = torch.load(checkpoint_name)
             model = set_peft_model_state_dict(model, adapters_weights)
         else:
-            print(f"Checkpoint {checkpoint_name} not found")
+            raise ValueError(f"Checkpoint {checkpoint_name} not found")
 
     # Be more transparent about the % of trainable params.
     model.print_trainable_parameters()
@@ -197,15 +265,15 @@ def train(
             optim="adamw_torch",
             evaluation_strategy="steps" if val_set_size > 0 else "no",
             save_strategy="steps",
-            eval_steps=200 if val_set_size > 0 else None,
+            eval_steps=save_steps if val_set_size > 0 else None,
             save_steps=save_steps,
             output_dir=output_dir,
             save_total_limit=save_total_limit,
             load_best_model_at_end=True if val_set_size > 0 else False,
             ddp_find_unused_parameters=False if ddp else None,
             group_by_length=group_by_length,
-            # report_to="wandb" if use_wandb else None,
-            # run_name=wandb_run_name if use_wandb else None,
+            report_to="wandb" if use_wandb else None,
+            run_name=wandb_run_name if use_wandb else None,
         ),
         data_collator=transformers.DataCollatorForSeq2Seq(
             tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
@@ -217,24 +285,16 @@ def train(
         os.makedirs(output_dir)
     with open(os.path.join(output_dir, "trainer_args.json"), 'w') as trainer_args_json_file:
         json.dump(trainer.args.to_dict(), trainer_args_json_file, indent=2)
-    with open(os.path.join(output_dir, "finetune_params.json"), 'w') as finetune_params_json_file:
-        finetune_params = {
-            'micro_batch_size': micro_batch_size,
-            'gradient_accumulation_steps': gradient_accumulation_steps,
-            'num_train_epochs': num_train_epochs,
-            'learning_rate': learning_rate,
-            'cutoff_len': cutoff_len,
-            'lora_r': lora_r,
-            'lora_alpha': lora_alpha,
-            'lora_dropout': lora_dropout,
-            'lora_target_modules': lora_target_modules,
-            'train_on_inputs': train_on_inputs,
-            'group_by_length': group_by_length,
-            'save_steps': save_steps,
-            'save_total_limit': save_total_limit,
-            'logging_steps': logging_steps,
-        }
-        json.dump(finetune_params, finetune_params_json_file, indent=2)
+    with open(os.path.join(output_dir, "finetune_args.json"), 'w') as finetune_args_json_file:
+        json.dump(finetune_args, finetune_args_json_file, indent=2)
+
+    # Not working, will only give us ["prompt", "completion", "input_ids", "attention_mask", "labels"]
+    # if train_data:
+    #     with open(os.path.join(output_dir, "train_dataset_samples.json"), 'w') as file:
+    #         json.dump(list(train_data[:100]), file, indent=2)
+    # if val_data:
+    #     with open(os.path.join(output_dir, "eval_dataset_samples.json"), 'w') as file:
+    #         json.dump(list(val_data[:100]), file, indent=2)
 
     model.config.use_cache = False
 
@@ -261,4 +321,7 @@ def train(
     with open(os.path.join(output_dir, "train_output.json"), 'w') as train_output_json_file:
         json.dump(train_output, train_output_json_file, indent=2)
 
+    if use_wandb and wandb:
+        wandb.finish()
+
     return train_output

llama_lora/lib/get_device.py

@@ -0,0 +1,15 @@
+import torch
+
+
+def get_device():
+    device ="cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+
+    try:
+        if torch.backends.mps.is_available():
+            device = "mps"
+    except:  # noqa: E722
+        pass
+
+    return device

llama_lora/lib/inference.py

@@ -0,0 +1,86 @@
+import torch
+import transformers
+
+from .get_device import get_device
+from .streaming_generation_utils import Iteratorize, Stream
+
+def generate(
+    # model
+    model,
+    tokenizer,
+    # input
+    prompt,
+    generation_config,
+    max_new_tokens,
+    stopping_criteria=[],
+    # output options
+    stream_output=False
+):
+    device = get_device()
+
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(device)
+    generate_params = {
+        "input_ids": input_ids,
+        "generation_config": generation_config,
+        "return_dict_in_generate": True,
+        "output_scores": True,
+        "max_new_tokens": max_new_tokens,
+        "stopping_criteria": transformers.StoppingCriteriaList() + stopping_criteria
+    }
+
+    skip_special_tokens = True
+
+    if '/dolly' in tokenizer.name_or_path:
+        # dolly has additional_special_tokens as ['### End', '### Instruction:', '### Response:'], skipping them will break the prompter's reply extraction.
+        skip_special_tokens = False
+        # Ensure generation stops once it generates "### End"
+        end_key_token_id = tokenizer.encode("### End")
+        end_key_token_id = end_key_token_id[0]  # 50277
+        if isinstance(generate_params['generation_config'].eos_token_id, str):
+            generate_params['generation_config'].eos_token_id = [generate_params['generation_config'].eos_token_id]
+        elif not generate_params['generation_config'].eos_token_id:
+            generate_params['generation_config'].eos_token_id = []
+        generate_params['generation_config'].eos_token_id.append(end_key_token_id)
+
+    if stream_output:
+        # Stream the reply 1 token at a time.
+        # This is based on the trick of using 'stopping_criteria' to create an iterator,
+        # from https://github.com/oobabooga/text-generation-webui/blob/ad37f396fc8bcbab90e11ecf17c56c97bfbd4a9c/modules/text_generation.py#L216-L243.
+        generation_output = None
+
+        def generate_with_callback(callback=None, **kwargs):
+            nonlocal generation_output
+            kwargs["stopping_criteria"].insert(
+                0,
+                Stream(callback_func=callback)
+            )
+            with torch.no_grad():
+                generation_output = model.generate(**kwargs)
+
+        def generate_with_streaming(**kwargs):
+            return Iteratorize(
+                generate_with_callback, kwargs, callback=None
+            )
+
+        with generate_with_streaming(**generate_params) as generator:
+            for output in generator:
+                decoded_output = tokenizer.decode(output, skip_special_tokens=skip_special_tokens)
+                yield decoded_output, output
+                if output[-1] in [tokenizer.eos_token_id]:
+                    break
+
+        if generation_output:
+            output = generation_output.sequences[0]
+            decoded_output = tokenizer.decode(output, skip_special_tokens=skip_special_tokens)
+            yield decoded_output, output
+
+        return  # early return for stream_output
+
+    # Without streaming
+    with torch.no_grad():
+        generation_output = model.generate(**generate_params)
+    output = generation_output.sequences[0]
+    decoded_output = tokenizer.decode(output, skip_special_tokens=skip_special_tokens)
+    yield decoded_output, output
+    return

llama_lora/models.py

@@ -2,25 +2,14 @@ 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
-
-
-def get_device():
-    if torch.cuda.is_available():
-        return "cuda"
-    else:
-        return "cpu"
-
-    try:
-        if torch.backends.mps.is_available():
-            return "mps"
-    except:  # noqa: E722
-        pass
+from .lib.get_device import get_device
 
 
 def get_new_base_model(base_model_name):
@@ -41,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,
@@ -50,19 +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
         )
 
-    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
+    tokenizer = get_tokenizer(base_model_name)
+
+    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
 
@@ -75,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
@@ -148,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

@@ -1,6 +1,8 @@
 import os
 import json
 import time
+import traceback
+import re
 from datetime import datetime
 import gradio as gr
 import math
@@ -15,7 +17,8 @@ from ..models import (
 from ..utils.data import (
     get_available_template_names,
     get_available_dataset_names,
-    get_dataset_content
+    get_dataset_content,
+    get_available_lora_model_names
 )
 from ..utils.prompter import Prompter
 
@@ -47,13 +50,16 @@ def reload_selections(current_template, current_dataset):
     current_dataset = current_dataset or next(
         iter(available_dataset_names), None)
 
+    available_lora_models = ["-"] + get_available_lora_model_names()
+
     return (
         gr.Dropdown.update(
             choices=available_template_names_with_none,
             value=current_template),
         gr.Dropdown.update(
             choices=available_dataset_names,
-            value=current_dataset)
+            value=current_dataset),
+        gr.Dropdown.update(choices=available_lora_models)
     )
 
 
@@ -79,56 +85,47 @@ def load_sample_dataset_to_text_input(format):
         return gr.Code.update(value=sample_plain_text_value)
 
 
-def process_json_dataset(data, only_first_n_items=None):
-    if not isinstance(data, list):
-        raise ValueError("The dataset is not an array of objects.")
-
-    if only_first_n_items is not None:
-        data = data[:only_first_n_items]
-
-    first_item = get_val_from_arr(data, 0, None)
-
-    if first_item is None:
-        raise ValueError("The dataset is empty.")
-    if not isinstance(first_item, dict):
-        raise ValueError("The dataset is not an array of objects.")
-
-    # Convert OpenAI fine-tuning dataset to LLaMA LoRA style
-    if "completion" in first_item and "output" not in first_item:
-        data = [
-            {"output" if k == "completion" else k: v for k, v in d.items()}
-            for d in data]
-        first_item = get_val_from_arr(data, 0, None)
-
-    # Flatten Stanford Alpaca style instances
-    if "instances" in first_item and isinstance(first_item["instances"], list):
-        data = [
-            {"output" if k == "completion" else k: v for k, v in d.items()}
-            for d in data]
-        flattened_data = []
-        for item in data:
-            for instance in item["instances"]:
-                d = {k: v for k, v in item.items() if k != "instances"}
-                d.update(instance)
-                flattened_data.append(d)
-        data = flattened_data
-        first_item = get_val_from_arr(data, 0, None)
-
-    if "output" not in first_item:
-        raise ValueError(
-            "The data does not contains an \"output\" or \"completion\".")
-
-    # Put all variables under the "variables" key if it does not exists
-    if "variables" not in first_item:
-        data = [
-            {
-                "variables":
-                    {k: v for k, v in d.items() if k != "output"},
-                "output":
-                    d["output"]
-            }
-            for d in data
-        ]
+def get_data_from_input(load_dataset_from, dataset_text, dataset_text_format,
+                        dataset_plain_text_input_variables_separator,
+                        dataset_plain_text_input_and_output_separator,
+                        dataset_plain_text_data_separator,
+                        dataset_from_data_dir, prompter):
+    if load_dataset_from == "Text Input":
+        if dataset_text_format == "JSON":
+            data = json.loads(dataset_text)
+
+        elif dataset_text_format == "JSON Lines":
+            lines = dataset_text.split('\n')
+            data = []
+            for i, line in enumerate(lines):
+                line_number = i + 1
+                try:
+                    data.append(json.loads(line))
+                except Exception as e:
+                    raise ValueError(
+                        f"Error parsing JSON on line {line_number}: {e}")
+
+        else:  # Plain Text
+            data = parse_plain_text_input(
+                dataset_text,
+                (
+                    dataset_plain_text_input_variables_separator or
+                    default_dataset_plain_text_input_variables_separator
+                ).replace("\\n", "\n"),
+                (
+                    dataset_plain_text_input_and_output_separator or
+                    default_dataset_plain_text_input_and_output_separator
+                ).replace("\\n", "\n"),
+                (
+                    dataset_plain_text_data_separator or
+                    default_dataset_plain_text_data_separator
+                ).replace("\\n", "\n"),
+                prompter.get_variable_names()
+            )
+
+    else:  # Load dataset from data directory
+        data = get_dataset_content(dataset_from_data_dir)
+
     return data
 
 
@@ -141,72 +138,92 @@ def refresh_preview(
     dataset_plain_text_input_variables_separator,
     dataset_plain_text_input_and_output_separator,
     dataset_plain_text_data_separator,
-    preview_show_actual_prompt,
+    max_preview_count,
 ):
     try:
-        max_preview_count = 100
         prompter = Prompter(template)
         variable_names = prompter.get_variable_names()
 
-        if load_dataset_from == "Text Input":
-            if dataset_text_format == "JSON":
-                data = json.loads(dataset_text)
-                data = process_json_dataset(data)
-
-            elif dataset_text_format == "JSON Lines":
-                lines = dataset_text.split('\n')
-                data = []
-                for i, line in enumerate(lines):
-                    line_number = i + 1
-                    try:
-                        data.append(json.loads(line))
-                    except Exception as e:
-                        raise ValueError(
-                            f"Error parsing JSON on line {line_number}: {e}")
-
-                data = process_json_dataset(data)
-
-            else:  # Plain Text
-                data = parse_plain_text_input(
-                    dataset_text,
-                    (
-                        dataset_plain_text_input_variables_separator or
-                        default_dataset_plain_text_input_variables_separator
-                    ).replace("\\n", "\n"),
-                    (
-                        dataset_plain_text_input_and_output_separator or
-                        default_dataset_plain_text_input_and_output_separator
-                    ).replace("\\n", "\n"),
-                    (
-                        dataset_plain_text_data_separator or
-                        default_dataset_plain_text_data_separator
-                    ).replace("\\n", "\n"),
-                    variable_names
-                )
+        data = get_data_from_input(
+            load_dataset_from=load_dataset_from,
+            dataset_text=dataset_text,
+            dataset_text_format=dataset_text_format,
+            dataset_plain_text_input_variables_separator=dataset_plain_text_input_variables_separator,
+            dataset_plain_text_input_and_output_separator=dataset_plain_text_input_and_output_separator,
+            dataset_plain_text_data_separator=dataset_plain_text_data_separator,
+            dataset_from_data_dir=dataset_from_data_dir,
+            prompter=prompter
+        )
+
+        train_data = prompter.get_train_data_from_dataset(
+            data, max_preview_count)
 
-        else:  # Load dataset from data directory
-            data = get_dataset_content(dataset_from_data_dir)
-            data = process_json_dataset(data)
+        train_data = train_data[:max_preview_count]
 
         data_count = len(data)
-        headers = variable_names
+
+        headers = ['Prompt', 'Completion']
         preview_data = [
-            [item['variables'].get(name, "") for name in variable_names]
-            for item in data[:max_preview_count]
+            [item.get("prompt", ""), item.get("completion", "")]
+            for item in train_data
         ]
 
-        if preview_show_actual_prompt:
-            headers = headers + ["Prompt (actual input)"]
-            rendered = [prompter.generate_prompt(
-                item['variables']) for item in data[:max_preview_count]]
-            preview_data = result = [d + [i]
-                                     for d, i in zip(preview_data, rendered)]
+        if not prompter.template_module:
+            variable_names = prompter.get_variable_names()
+            headers += [f"Variable: {variable_name}" for variable_name in variable_names]
+            variables = [
+                [item.get(f"_var_{name}", "") for name in variable_names]
+                for item in train_data
+            ]
+            preview_data = [d + v for d, v in zip(preview_data, variables)]
+
+        preview_info_message = f"The dataset has about {data_count} item(s)."
+        if data_count > max_preview_count:
+            preview_info_message += f" Previewing the first {max_preview_count}."
+
+        info_message = f"about {data_count} item(s)."
+        if load_dataset_from == "Data Dir":
+            info_message = "This dataset contains about " + info_message
+        update_message = gr.Markdown.update(info_message, visible=True)
 
-        headers = headers + ["Completion (output)"]
-        preview_data = result = [pd + [d['output']]
-                                 for pd, d in zip(preview_data, data[:max_preview_count])]
+        return gr.Dataframe.update(value={'data': preview_data, 'headers': headers}), gr.Markdown.update(preview_info_message), update_message, update_message
+    except Exception as e:
+        update_message = gr.Markdown.update(
+            f"<span class=\"finetune_dataset_error_message\">Error: {e}.</span>", visible=True)
+        return gr.Dataframe.update(value={'data': [], 'headers': []}), gr.Markdown.update("Set the dataset in the \"Prepare\" tab, then preview it here."), update_message, update_message
 
-        preview_info_message = f"The dataset has a total of {data_count} item(s)."
+
+def refresh_dataset_items_count(
+    template,
+    load_dataset_from,
+    dataset_from_data_dir,
+    dataset_text,
+    dataset_text_format,
+    dataset_plain_text_input_variables_separator,
+    dataset_plain_text_input_and_output_separator,
+    dataset_plain_text_data_separator,
+    max_preview_count,
+):
+    try:
+        prompter = Prompter(template)
+        variable_names = prompter.get_variable_names()
+
+        data = get_data_from_input(
+            load_dataset_from=load_dataset_from,
+            dataset_text=dataset_text,
+            dataset_text_format=dataset_text_format,
+            dataset_plain_text_input_variables_separator=dataset_plain_text_input_variables_separator,
+            dataset_plain_text_input_and_output_separator=dataset_plain_text_input_and_output_separator,
+            dataset_plain_text_data_separator=dataset_plain_text_data_separator,
+            dataset_from_data_dir=dataset_from_data_dir,
+            prompter=prompter
+        )
+
+        train_data = prompter.get_train_data_from_dataset(
+            data)
+        data_count = len(train_data)
+
+        preview_info_message = f"The dataset contains {data_count} item(s)."
         if data_count > max_preview_count:
             preview_info_message += f" Previewing the first {max_preview_count}."
 
@@ -215,11 +232,22 @@ def refresh_preview(
             info_message = "This dataset contains " + info_message
         update_message = gr.Markdown.update(info_message, visible=True)
 
-        return gr.Dataframe.update(value={'data': preview_data, 'headers': headers}), gr.Markdown.update(preview_info_message), update_message, update_message
+        return gr.Markdown.update(preview_info_message), update_message, update_message, gr.Slider.update(maximum=math.floor(data_count / 2))
     except Exception as e:
         update_message = gr.Markdown.update(
             f"<span class=\"finetune_dataset_error_message\">Error: {e}.</span>", visible=True)
-        return gr.Dataframe.update(value={'data': [], 'headers': []}), gr.Markdown.update("Set the dataset in the \"Prepare\" tab, then preview it here."), update_message, update_message
+
+        trace = traceback.format_exc()
+        traces = [s.strip() for s in re.split("\n * File ", trace)]
+        templates_path = os.path.join(Global.data_dir, "templates")
+        traces_to_show = [s for s in traces if os.path.join(
+            Global.data_dir, "templates") in s]
+        traces_to_show = [re.sub(" *\n *", ": ", s) for s in traces_to_show]
+        if len(traces_to_show) > 0:
+            update_message = gr.Markdown.update(
+                f"<span class=\"finetune_dataset_error_message\">Error: {e} ({','.join(traces_to_show)}).</span>", visible=True)
+
+        return gr.Markdown.update("Set the dataset in the \"Prepare\" tab, then preview it here."), update_message, update_message, gr.Slider.update(maximum=1)
 
 
 def parse_plain_text_input(
@@ -258,7 +286,7 @@ def do_train(
     dataset_plain_text_data_separator,
     # Training Options
     max_seq_length,
-    evaluate_data_percentage,
+    evaluate_data_count,
     micro_batch_size,
     gradient_accumulation_steps,
     epochs,
@@ -268,14 +296,27 @@ def do_train(
     lora_alpha,
     lora_dropout,
     lora_target_modules,
-    model_name,
     save_steps,
     save_total_limit,
     logging_steps,
+    model_name,
+    continue_from_model,
+    continue_from_checkpoint,
     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
+
+        resume_from_checkpoint = None
+        if continue_from_model == "-" or continue_from_model == "None":
+            continue_from_model = None
+        if continue_from_checkpoint == "-" or continue_from_checkpoint == "None":
+            continue_from_checkpoint = None
+        if continue_from_model:
+            resume_from_checkpoint = os.path.join(Global.data_dir, "lora_models", continue_from_model)
+            if continue_from_checkpoint:
+                resume_from_checkpoint = os.path.join(resume_from_checkpoint, continue_from_checkpoint)
+
         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')):
@@ -288,56 +329,22 @@ def do_train(
         unload_models()  # Need RAM for training
 
         prompter = Prompter(template)
-        variable_names = prompter.get_variable_names()
-
-        if load_dataset_from == "Text Input":
-            if dataset_text_format == "JSON":
-                data = json.loads(dataset_text)
-                data = process_json_dataset(data)
-
-            elif dataset_text_format == "JSON Lines":
-                lines = dataset_text.split('\n')
-                data = []
-                for i, line in enumerate(lines):
-                    line_number = i + 1
-                    try:
-                        data.append(json.loads(line))
-                    except Exception as e:
-                        raise ValueError(
-                            f"Error parsing JSON on line {line_number}: {e}")
-
-                data = process_json_dataset(data)
-
-            else:  # Plain Text
-                data = parse_plain_text_input(
-                    dataset_text,
-                    (
-                        dataset_plain_text_input_variables_separator or
-                        default_dataset_plain_text_input_variables_separator
-                    ).replace("\\n", "\n"),
-                    (
-                        dataset_plain_text_input_and_output_separator or
-                        default_dataset_plain_text_input_and_output_separator
-                    ).replace("\\n", "\n"),
-                    (
-                        dataset_plain_text_data_separator or
-                        default_dataset_plain_text_data_separator
-                    ).replace("\\n", "\n"),
-                    variable_names
-                )
-
-        else:  # Load dataset from data directory
-            data = get_dataset_content(dataset_from_data_dir)
-            data = process_json_dataset(data)
+        # variable_names = prompter.get_variable_names()
+
+        data = get_data_from_input(
+            load_dataset_from=load_dataset_from,
+            dataset_text=dataset_text,
+            dataset_text_format=dataset_text_format,
+            dataset_plain_text_input_variables_separator=dataset_plain_text_input_variables_separator,
+            dataset_plain_text_input_and_output_separator=dataset_plain_text_input_and_output_separator,
+            dataset_plain_text_data_separator=dataset_plain_text_data_separator,
+            dataset_from_data_dir=dataset_from_data_dir,
+            prompter=prompter
+        )
 
-        data_count = len(data)
-        evaluate_data_count = math.ceil(data_count * evaluate_data_percentage)
+        train_data = prompter.get_train_data_from_dataset(data)
 
-        train_data = [
-            {
-                'prompt': prompter.generate_prompt(d['variables']),
-                'completion': d['output']}
-            for d in data]
+        data_count = len(train_data)
 
         def get_progress_text(epoch, epochs, last_loss):
             progress_detail = f"Epoch {math.ceil(epoch)}/{epochs}"
@@ -380,6 +387,8 @@ Train options: {json.dumps({
     'lora_dropout': lora_dropout,
     'lora_target_modules': lora_target_modules,
     'model_name': model_name,
+    'continue_from_model': continue_from_model,
+    'continue_from_checkpoint': continue_from_checkpoint,
 }, indent=2)}
 
 Train data (first 10):
@@ -390,7 +399,7 @@ Train data (first 10):
             return message
 
         if not should_training_progress_track_tqdm:
-            progress(0, desc="Preparing model for training...")
+            progress(0, desc=f"Preparing model {base_model_name} for training...")
 
         log_history = []
 
@@ -449,26 +458,37 @@ Train data (first 10):
                 'dataset_rows': len(train_data),
                 'timestamp': time.time(),
 
-                'max_seq_length': max_seq_length,
-                'train_on_inputs': train_on_inputs,
+                # These will be saved in another JSON file by the train function
+                # 'max_seq_length': max_seq_length,
+                # 'train_on_inputs': train_on_inputs,
 
-                'micro_batch_size': micro_batch_size,
-                'gradient_accumulation_steps': gradient_accumulation_steps,
-                'epochs': epochs,
-                'learning_rate': learning_rate,
+                # 'micro_batch_size': micro_batch_size,
+                # 'gradient_accumulation_steps': gradient_accumulation_steps,
+                # 'epochs': epochs,
+                # 'learning_rate': learning_rate,
 
-                'evaluate_data_percentage': evaluate_data_percentage,
+                # 'evaluate_data_count': evaluate_data_count,
 
-                'lora_r': lora_r,
-                'lora_alpha': lora_alpha,
-                'lora_dropout': lora_dropout,
-                'lora_target_modules': lora_target_modules,
+                # 'lora_r': lora_r,
+                # 'lora_alpha': lora_alpha,
+                # 'lora_dropout': lora_dropout,
+                # 'lora_target_modules': lora_target_modules,
             }
+            if continue_from_model:
+                info['continued_from_model'] = continue_from_model
+                if continue_from_checkpoint:
+                    info['continued_from_checkpoint'] = continue_from_checkpoint
             json.dump(info, info_json_file, indent=2)
 
         if not should_training_progress_track_tqdm:
             progress(0, desc="Train starting...")
 
+        wandb_group = template
+        wandb_tags = [f"template:{template}"]
+        if load_dataset_from == "Data Dir" and dataset_from_data_dir:
+            wandb_group += f"/{dataset_from_data_dir}"
+            wandb_tags.append(f"dataset:{dataset_from_data_dir}")
+
         train_output = Global.train_fn(
             base_model,  # base_model
             tokenizer,  # tokenizer
@@ -487,11 +507,16 @@ Train data (first 10):
             lora_target_modules,  # lora_target_modules
             train_on_inputs,  # train_on_inputs
             False,  # group_by_length
-            None,  # resume_from_checkpoint
+            resume_from_checkpoint,  # resume_from_checkpoint
             save_steps,  # save_steps
             save_total_limit,  # save_total_limit
             logging_steps,  # logging_steps
-            training_callbacks  # callbacks
+            training_callbacks,  # callbacks
+            Global.wandb_api_key,  # wandb_api_key
+            Global.default_wandb_project if Global.enable_wandb else None,  # wandb_project
+            wandb_group,  # wandb_group
+            model_name,  # wandb_run_name
+            wandb_tags  # wandb_tags
         )
 
         logs_str = "\n".join([json.dumps(log)
@@ -515,6 +540,146 @@ def do_abort_training():
     Global.should_stop_training = True
 
 
+def handle_continue_from_model_change(model_name):
+    try:
+        lora_models_directory_path = os.path.join(
+            Global.data_dir, "lora_models")
+        lora_model_directory_path = os.path.join(
+            lora_models_directory_path, model_name)
+        all_files = os.listdir(lora_model_directory_path)
+        checkpoints = [
+            file for file in all_files if file.startswith("checkpoint-")]
+        checkpoints = ["-"] + checkpoints
+        can_load_params = "finetune_params.json" in all_files or "finetune_args.json" in all_files
+        return gr.Dropdown.update(choices=checkpoints, value="-"), gr.Button.update(visible=can_load_params), gr.Markdown.update(value="", visible=False)
+    except Exception:
+        pass
+    return gr.Dropdown.update(choices=["-"], value="-"), gr.Button.update(visible=False), gr.Markdown.update(value="", visible=False)
+
+
+def handle_load_params_from_model(
+    model_name,
+    max_seq_length,
+    evaluate_data_count,
+    micro_batch_size,
+    gradient_accumulation_steps,
+    epochs,
+    learning_rate,
+    train_on_inputs,
+    lora_r,
+    lora_alpha,
+    lora_dropout,
+    lora_target_modules,
+    save_steps,
+    save_total_limit,
+    logging_steps,
+    lora_target_module_choices,
+):
+    error_message = ""
+    notice_message = ""
+    unknown_keys = []
+    try:
+        lora_models_directory_path = os.path.join(
+            Global.data_dir, "lora_models")
+        lora_model_directory_path = os.path.join(
+            lora_models_directory_path, model_name)
+
+        data = {}
+        possible_files = ["finetune_params.json", "finetune_args.json"]
+        for file in possible_files:
+            try:
+                with open(os.path.join(lora_model_directory_path, file), "r") as f:
+                    data = json.load(f)
+            except FileNotFoundError:
+                pass
+
+        for key, value in data.items():
+            if key == "max_seq_length":
+                max_seq_length = value
+            if key == "cutoff_len":
+                cutoff_len = value
+            elif key == "evaluate_data_count":
+                evaluate_data_count = value
+            elif key == "val_set_size":
+                evaluate_data_count = value
+            elif key == "micro_batch_size":
+                micro_batch_size = value
+            elif key == "gradient_accumulation_steps":
+                gradient_accumulation_steps = value
+            elif key == "epochs":
+                epochs = value
+            elif key == "num_train_epochs":
+                epochs = value
+            elif key == "learning_rate":
+                learning_rate = value
+            elif key == "train_on_inputs":
+                train_on_inputs = value
+            elif key == "lora_r":
+                lora_r = value
+            elif key == "lora_alpha":
+                lora_alpha = value
+            elif key == "lora_dropout":
+                lora_dropout = value
+            elif key == "lora_target_modules":
+                lora_target_modules = value
+                for element in value:
+                    if element not in lora_target_module_choices:
+                        lora_target_module_choices.append(element)
+            elif key == "save_steps":
+                save_steps = value
+            elif key == "save_total_limit":
+                save_total_limit = value
+            elif key == "logging_steps":
+                logging_steps = value
+            elif key == "group_by_length":
+                pass
+            elif key == "resume_from_checkpoint":
+                pass
+            else:
+                unknown_keys.append(key)
+    except Exception as e:
+        error_message = str(e)
+
+    if len(unknown_keys) > 0:
+        notice_message = f"Note: cannot restore unknown arg: {', '.join([f'`{x}`' for x in unknown_keys])}"
+
+    message = ". ".join([x for x in [error_message, notice_message] if x])
+
+    has_message = False
+    if message:
+        message += "."
+        has_message = True
+
+    return (
+        gr.Markdown.update(value=message, visible=has_message),
+        max_seq_length,
+        evaluate_data_count,
+        micro_batch_size,
+        gradient_accumulation_steps,
+        epochs,
+        learning_rate,
+        train_on_inputs,
+        lora_r,
+        lora_alpha,
+        lora_dropout,
+        gr.CheckboxGroup.update(value=lora_target_modules, choices=lora_target_module_choices),
+        save_steps,
+        save_total_limit,
+        logging_steps,
+        lora_target_module_choices,
+    )
+
+
+default_lora_target_module_choices = ["q_proj", "k_proj", "v_proj", "o_proj"]
+
+
+def handle_lora_target_modules_add(choices, new_module, selected_modules):
+    choices.append(new_module)
+    selected_modules.append(new_module)
+
+    return (choices, "", gr.CheckboxGroup.update(value=selected_modules, choices=choices))
+
+
 def finetune_ui():
     things_that_might_timeout = []
 
@@ -606,9 +771,13 @@ def finetune_ui():
                         "Set the dataset in the \"Prepare\" tab, then preview it here.",
                         elem_id="finetune_dataset_preview_info_message"
                     )
-                    finetune_dataset_preview_show_actual_prompt = gr.Checkbox(
-                        label="Show actual prompt",
-                        elem_id="finetune_dataset_preview_show_actual_prompt"
+                    finetune_dataset_preview_count = gr.Number(
+                        label="Preview items count",
+                        value=10,
+                        # minimum=1,
+                        # maximum=100,
+                        precision=0,
+                        elem_id="finetune_dataset_preview_count"
                     )
                 finetune_dataset_preview = gr.Dataframe(
                     wrap=True, elem_id="finetune_dataset_preview")
@@ -633,25 +802,7 @@ def finetune_ui():
                 dataset_plain_text_data_separator,
             ]
             dataset_preview_inputs = dataset_inputs + \
-                [finetune_dataset_preview_show_actual_prompt]
-            for i in dataset_preview_inputs:
-                things_that_might_timeout.append(
-                    i.change(
-                        fn=refresh_preview,
-                        inputs=dataset_preview_inputs,
-                        outputs=[finetune_dataset_preview,
-                                 finetune_dataset_preview_info_message,
-                                 dataset_from_text_message,
-                                 dataset_from_data_dir_message
-                                 ]
-                    ))
-
-            things_that_might_timeout.append(reload_selections_button.click(
-                reload_selections,
-                inputs=[template, dataset_from_data_dir],
-                outputs=[template, dataset_from_data_dir],
-            )
-            )
+                [finetune_dataset_preview_count]
 
             with gr.Row():
                 max_seq_length = gr.Slider(
@@ -704,12 +855,43 @@ def finetune_ui():
                     info="The initial learning rate for the optimizer. A higher learning rate may speed up convergence but also cause instability or divergence. A lower learning rate may require more steps to reach optimal performance but also avoid overshooting or oscillating around local minima."
                 )
 
-                evaluate_data_percentage = gr.Slider(
-                    minimum=0, maximum=0.5, step=0.001, value=0,
-                    label="Evaluation Data Percentage",
-                    info="The percentage of data to be used for evaluation. This percentage of data will not be used for training and will be used to assess the performance of the model during the process."
+                evaluate_data_count = gr.Slider(
+                    minimum=0, maximum=1, step=1, value=0,
+                    label="Evaluation Data Count",
+                    info="The number of data to be used for evaluation. This amount of data will not be used for training and will be used to assess the performance of the model during the process."
                 )
 
+                with gr.Box(elem_id="finetune_continue_from_model_box"):
+                    with gr.Row():
+                        continue_from_model = gr.Dropdown(
+                            value="-",
+                            label="Continue from Model",
+                            choices=["-"],
+                            elem_id="finetune_continue_from_model"
+                        )
+                        continue_from_checkpoint = gr.Dropdown(
+                            value="-", label="Checkpoint", choices=["-"])
+                    with gr.Column():
+                        load_params_from_model_btn = gr.Button(
+                            "Load training parameters from selected model", visible=False)
+                        load_params_from_model_btn.style(
+                            full_width=False,
+                            size="sm")
+                        load_params_from_model_message = gr.Markdown(
+                            "", visible=False)
+
+                    things_that_might_timeout.append(
+                        continue_from_model.change(
+                            fn=handle_continue_from_model_change,
+                            inputs=[continue_from_model],
+                            outputs=[
+                                continue_from_checkpoint,
+                                load_params_from_model_btn,
+                                load_params_from_model_message
+                            ]
+                        )
+                    )
+
             with gr.Column():
                 lora_r = gr.Slider(
                     minimum=1, maximum=16, step=1, value=8,
@@ -729,12 +911,31 @@ def finetune_ui():
                     info="The dropout probability for LoRA, which controls the fraction of LoRA parameters that are set to zero during training. A larger lora_dropout increases the regularization effect of LoRA but also increases the risk of underfitting."
                 )
 
+                lora_target_module_choices = gr.State(value=default_lora_target_module_choices)
+
                 lora_target_modules = gr.CheckboxGroup(
                     label="LoRA Target Modules",
-                    choices=["q_proj", "k_proj", "v_proj", "o_proj"],
+                    choices=default_lora_target_module_choices,
                     value=["q_proj", "v_proj"],
-                    info="Modules to replace with LoRA."
+                    info="Modules to replace with LoRA.",
+                    elem_id="finetune_lora_target_modules"
                 )
+                with gr.Box(elem_id="finetune_lora_target_modules_add_box"):
+                    with gr.Row():
+                        lora_target_modules_add = gr.Textbox(
+                            lines=1, max_lines=1, show_label=False,
+                            elem_id="finetune_lora_target_modules_add"
+                        )
+                        lora_target_modules_add_btn = gr.Button(
+                            "Add",
+                            elem_id="finetune_lora_target_modules_add_btn"
+                        )
+                        lora_target_modules_add_btn.style(full_width=False, size="sm")
+                things_that_might_timeout.append(lora_target_modules_add_btn.click(
+                    handle_lora_target_modules_add,
+                    inputs=[lora_target_module_choices, lora_target_modules_add, lora_target_modules],
+                    outputs=[lora_target_module_choices, lora_target_modules_add, lora_target_modules],
+                ))
 
                 with gr.Row():
                     logging_steps = gr.Number(
@@ -759,6 +960,7 @@ def finetune_ui():
                 with gr.Column():
                     model_name = gr.Textbox(
                         lines=1, label="LoRA Model Name", value=random_name,
+                        max_lines=1,
                         info="The name of the new LoRA model.",
                         elem_id="finetune_model_name",
                     )
@@ -778,6 +980,59 @@ def finetune_ui():
                             elem_id="finetune_confirm_stop_btn"
                         )
 
+        things_that_might_timeout.append(reload_selections_button.click(
+            reload_selections,
+            inputs=[template, dataset_from_data_dir],
+            outputs=[template, dataset_from_data_dir, continue_from_model],
+        ))
+
+        for i in dataset_preview_inputs:
+            things_that_might_timeout.append(
+                i.change(
+                    fn=refresh_preview,
+                    inputs=dataset_preview_inputs,
+                    outputs=[
+                        finetune_dataset_preview,
+                        finetune_dataset_preview_info_message,
+                        dataset_from_text_message,
+                        dataset_from_data_dir_message
+                    ]
+                ).then(
+                    fn=refresh_dataset_items_count,
+                    inputs=dataset_preview_inputs,
+                    outputs=[
+                        finetune_dataset_preview_info_message,
+                        dataset_from_text_message,
+                        dataset_from_data_dir_message,
+                        evaluate_data_count,
+                    ]
+                ))
+
+        finetune_args = [
+            max_seq_length,
+            evaluate_data_count,
+            micro_batch_size,
+            gradient_accumulation_steps,
+            epochs,
+            learning_rate,
+            train_on_inputs,
+            lora_r,
+            lora_alpha,
+            lora_dropout,
+            lora_target_modules,
+            save_steps,
+            save_total_limit,
+            logging_steps,
+        ]
+
+        things_that_might_timeout.append(
+            load_params_from_model_btn.click(
+                fn=handle_load_params_from_model,
+                inputs=[continue_from_model] + finetune_args + [lora_target_module_choices],
+                outputs=[load_params_from_model_message] + finetune_args + [lora_target_module_choices]
+            )
+        )
+
         train_output = gr.Text(
             "Training results will be shown here.",
             label="Train Output",
@@ -785,22 +1040,10 @@ def finetune_ui():
 
         train_progress = train_btn.click(
             fn=do_train,
-            inputs=(dataset_inputs + [
-                max_seq_length,
-                evaluate_data_percentage,
-                micro_batch_size,
-                gradient_accumulation_steps,
-                epochs,
-                learning_rate,
-                train_on_inputs,
-                lora_r,
-                lora_alpha,
-                lora_dropout,
-                lora_target_modules,
+            inputs=(dataset_inputs + finetune_args + [
                 model_name,
-                save_steps,
-                save_total_limit,
-                logging_steps,
+                continue_from_model,
+                continue_from_checkpoint,
             ]),
             outputs=train_output
         )

llama_lora/ui/inference_ui.py

@@ -8,12 +8,12 @@ from transformers import GenerationConfig
 
 from ..globals import Global
 from ..models import get_model, get_tokenizer, get_device
+from ..lib.inference import generate
 from ..utils.data import (
     get_available_template_names,
     get_available_lora_model_names,
     get_info_of_available_lora_model)
 from ..utils.prompter import Prompter
-from ..utils.callbacks import Iteratorize, Stream
 
 device = get_device()
 
@@ -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:
@@ -103,8 +103,6 @@ def do_inference(
         tokenizer = get_tokenizer(base_model_name)
         model = get_model(base_model_name, lora_model_name)
 
-        inputs = tokenizer(prompt, return_tensors="pt")
-        input_ids = inputs["input_ids"].to(device)
         generation_config = GenerationConfig(
             temperature=temperature,
             top_p=top_p,
@@ -113,103 +111,55 @@ def do_inference(
             num_beams=num_beams,
         )
 
-        generate_params = {
-            "input_ids": input_ids,
-            "generation_config": generation_config,
-            "return_dict_in_generate": True,
-            "output_scores": True,
-            "max_new_tokens": max_new_tokens,
-        }
-
         def ui_generation_stopping_criteria(input_ids, score, **kwargs):
             if Global.should_stop_generating:
                 return True
             return False
 
         Global.should_stop_generating = False
-        generate_params.setdefault(
-            "stopping_criteria", transformers.StoppingCriteriaList()
-        )
-        generate_params["stopping_criteria"].append(
-            ui_generation_stopping_criteria
-        )
-
-        if stream_output:
-            # Stream the reply 1 token at a time.
-            # This is based on the trick of using 'stopping_criteria' to create an iterator,
-            # from https://github.com/oobabooga/text-generation-webui/blob/ad37f396fc8bcbab90e11ecf17c56c97bfbd4a9c/modules/text_generation.py#L216-L243.
-
-            def generate_with_callback(callback=None, **kwargs):
-                kwargs.setdefault(
-                    "stopping_criteria", transformers.StoppingCriteriaList()
-                )
-                kwargs["stopping_criteria"].append(
-                    Stream(callback_func=callback)
-                )
-                with torch.no_grad():
-                    model.generate(**kwargs)
-
-            def generate_with_streaming(**kwargs):
-                return Iteratorize(
-                    generate_with_callback, kwargs, callback=None
-                )
 
-            with generate_with_streaming(**generate_params) as generator:
-                for output in generator:
-                    # new_tokens = len(output) - len(input_ids[0])
-                    decoded_output = tokenizer.decode(output)
-
-                    if output[-1] in [tokenizer.eos_token_id]:
-                        break
-
-                    raw_output = None
-                    if show_raw:
-                        raw_output = str(output)
-                    response = prompter.get_response(decoded_output)
+        generation_args = {
+            'model': model,
+            'tokenizer': tokenizer,
+            'prompt': prompt,
+            'generation_config': generation_config,
+            'max_new_tokens': max_new_tokens,
+            'stopping_criteria': [ui_generation_stopping_criteria],
+            'stream_output': stream_output
+        }
 
-                    if Global.should_stop_generating:
-                        return
+        for (decoded_output, output) in generate(**generation_args):
+            raw_output_str = None
+            if show_raw:
+                raw_output_str = str(output)
+            response = prompter.get_response(decoded_output)
 
-                    yield (
-                        gr.Textbox.update(
-                            value=response, lines=inference_output_lines),
-                        raw_output)
-
-                    if Global.should_stop_generating:
-                        # If the user stops the generation, and then clicks the
-                        # generation button again, they may mysteriously landed
-                        # here, in the previous, should-be-stopped generation
-                        # function call, with the new generation function not be
-                        # called at all. To workaround this, we yield a message
-                        # and setting lines=1, and if the front-end JS detects
-                        # that lines has been set to 1 (rows="1" in HTML),
-                        # it will automatically click the generate button again
-                        # (gr.Textbox.update() does not support updating
-                        # elem_classes or elem_id).
-                        # [WORKAROUND-UI01]
-                        yield (
-                            gr.Textbox.update(
-                                value="Please retry", lines=1),
-                            None)
-            return  # early return for stream_output
-
-        # Without streaming
-        with torch.no_grad():
-            generation_output = model.generate(**generate_params)
-        s = generation_output.sequences[0]
-        output = tokenizer.decode(s)
-        raw_output = None
-        if show_raw:
-            raw_output = str(s)
-
-        response = prompter.get_response(output)
-        if Global.should_stop_generating:
-            return
+            if Global.should_stop_generating:
+                return
 
-        yield (
-            gr.Textbox.update(value=response, lines=inference_output_lines),
-            raw_output)
+            yield (
+                gr.Textbox.update(
+                    value=response, lines=inference_output_lines),
+                raw_output_str)
 
+            if Global.should_stop_generating:
+                # If the user stops the generation, and then clicks the
+                # generation button again, they may mysteriously landed
+                # here, in the previous, should-be-stopped generation
+                # function call, with the new generation function not be
+                # called at all. To workaround this, we yield a message
+                # and setting lines=1, and if the front-end JS detects
+                # that lines has been set to 1 (rows="1" in HTML),
+                # it will automatically click the generate button again
+                # (gr.Textbox.update() does not support updating
+                # elem_classes or elem_id).
+                # [WORKAROUND-UI01]
+                yield (
+                    gr.Textbox.update(
+                        value="Please retry", lines=1),
+                    None)
+
+        return
     except Exception as e:
         raise gr.Error(e)
 
@@ -229,7 +179,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"]
 
@@ -255,8 +205,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)
 
@@ -303,7 +257,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(
@@ -433,6 +387,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],
@@ -454,6 +410,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],
@@ -510,7 +474,7 @@ def inference_ui():
 
         // Workaround default value not shown.
         document.querySelector('#inference_lora_model input').value =
-          'tloen/alpaca-lora-7b';
+          'None';
       }, 100);
 
       // Add tooltips
@@ -654,6 +618,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,81 @@ 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;
+    }
+    #global_base_model_select ul {
+        z-index: 9999;
+        background: var(--block-background-fill);
+    }
+
+    #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 +252,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;
     }
@@ -413,6 +519,24 @@ def main_page_custom_css():
         margin: -32px -16px;
     }
 
+    #finetune_continue_from_model_box {
+        /* padding: 0; */
+    }
+    #finetune_continue_from_model_box .block {
+        border: 0;
+        box-shadow: none;
+        padding: 0;
+    }
+    #finetune_continue_from_model_box > * {
+        /* gap: 0; */
+    }
+    #finetune_continue_from_model_box button {
+        margin-top: 16px;
+    }
+    #finetune_continue_from_model {
+        flex-grow: 2;
+    }
+
     .finetune_dataset_error_message {
         color: var(--error-text-color) !important;
     }
@@ -428,10 +552,43 @@ def main_page_custom_css():
         white-space: pre-wrap;
     }
 
+    /*
+    #finetune_dataset_preview {
+        max-height: 100vh;
+        overflow: auto;
+        border: var(--block-border-width) solid var(--border-color-primary);
+        border-radius: var(--radius-lg);
+    }
+    #finetune_dataset_preview .table-wrap {
+        border: 0 !important;
+    }
+    */
+
     #finetune_max_seq_length {
         flex: 2;
     }
 
+    #finetune_lora_target_modules_add_box {
+        margin-top: -24px;
+        padding-top: 8px;
+        border-top-left-radius: 0;
+        border-top-right-radius: 0;
+        border-top: 0;
+    }
+    #finetune_lora_target_modules_add_box > * > .form {
+        border: 0;
+        box-shadow: none;
+    }
+    #finetune_lora_target_modules_add {
+        padding: 0;
+    }
+    #finetune_lora_target_modules_add input {
+        padding: 4px 8px;
+    }
+    #finetune_lora_target_modules_add_btn {
+        min-width: 60px;
+    }
+
     #finetune_save_total_limit,
     #finetune_save_steps,
     #finetune_logging_steps {
@@ -503,3 +660,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>
+        """

llama_lora/ui/tokenizer_ui.py

@@ -7,7 +7,7 @@ from ..models import get_tokenizer
 
 
 def handle_decode(encoded_tokens_json):
-    base_model_name = Global.default_base_model_name
+    base_model_name = Global.base_model_name
     try:
         encoded_tokens = json.loads(encoded_tokens_json)
         if Global.ui_dev_mode:
@@ -20,7 +20,7 @@ def handle_decode(encoded_tokens_json):
 
 
 def handle_encode(decoded_tokens):
-    base_model_name = Global.default_base_model_name
+    base_model_name = 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)

llama_lora/utils/data.py

@@ -30,19 +30,22 @@ def copy_sample_data_if_not_exists(source, destination):
 def get_available_template_names():
     templates_directory_path = os.path.join(Global.data_dir, "templates")
     all_files = os.listdir(templates_directory_path)
-    return [os.path.splitext(filename)[0] for filename in all_files if fnmatch.fnmatch(filename, "*.json")]
+    names = [filename.rstrip(".json") for filename in all_files if fnmatch.fnmatch(filename, "*.json") or fnmatch.fnmatch(filename, "*.py")]
+    return sorted(names)
 
 
 def get_available_dataset_names():
     datasets_directory_path = os.path.join(Global.data_dir, "datasets")
     all_files = os.listdir(datasets_directory_path)
-    return [filename for filename in all_files if fnmatch.fnmatch(filename, "*.json") or fnmatch.fnmatch(filename, "*.jsonl")]
+    names = [filename for filename in all_files if fnmatch.fnmatch(filename, "*.json") or fnmatch.fnmatch(filename, "*.jsonl")]
+    return sorted(names)
 
 
 def get_available_lora_model_names():
-    datasets_directory_path = os.path.join(Global.data_dir, "lora_models")
-    all_items = os.listdir(datasets_directory_path)
-    return [item for item in all_items if os.path.isdir(os.path.join(datasets_directory_path, item))]
+    lora_models_directory_path = os.path.join(Global.data_dir, "lora_models")
+    all_items = os.listdir(lora_models_directory_path)
+    names = [item for item in all_items if os.path.isdir(os.path.join(lora_models_directory_path, item))]
+    return sorted(names)
 
 
 def get_path_of_available_lora_model(name):

llama_lora/utils/prompter.py

@@ -5,13 +5,15 @@ From https://github.com/tloen/alpaca-lora/blob/main/utils/prompter.py
 
 import json
 import os.path as osp
+import importlib
+import itertools
 from typing import Union, List
 
 from ..globals import Global
 
 
 class Prompter(object):
-    __slots__ = ("template_name", "template", "_verbose")
+    __slots__ = ("template_name", "template", "template_module", "_verbose")
 
     def __init__(self, template_name: str = "", verbose: bool = False):
         self._verbose = verbose
@@ -21,12 +23,41 @@ class Prompter(object):
             self.template_name = "None"
             return
         self.template_name = template_name
+        self.template_module = None
 
-        file_name = osp.join(Global.data_dir, "templates",
-                             f"{template_name}.json")
-        if not osp.exists(file_name):
-            raise ValueError(f"Can't read {file_name}")
-        with open(file_name) as fp:
+        base_filename, ext = osp.splitext(template_name)
+        if ext == "":
+            filename = base_filename + ".json"
+        else:
+            filename = base_filename + ext
+
+        file_path = osp.join(Global.data_dir, "templates", filename)
+
+        if not osp.exists(file_path):
+            raise ValueError(f"Can't read {file_path}")
+
+        if ext == ".py":
+            template_module_spec = importlib.util.spec_from_file_location(
+                "template_module", file_path)
+            template_module = importlib.util.module_from_spec(
+                template_module_spec)
+            template_module_spec.loader.exec_module(template_module)
+            self.template_module = template_module
+
+            if not hasattr(template_module, "variables"):
+                raise ValueError(
+                    "The template module does not have a \"variables\" attribute.")
+
+            self.template = {
+                'variables': template_module.variables
+            }
+
+            if hasattr(template_module, "response_split"):
+                self.template["response_split"] = template_module.response_split
+
+            return
+
+        with open(file_path) as fp:
             self.template = json.load(fp)
         if self._verbose:
             print(
@@ -47,23 +78,31 @@ class Prompter(object):
                 res = variables.get("prompt", "")
         elif "variables" in self.template:
             variable_names = self.template.get("variables")
-            if type(variables) == dict:
-                variables = [variables.get(name, None)
-                             for name in variable_names]
-            if "default" not in self.template:
-                raise ValueError(
-                    f"The template {self.template_name} has \"variables\" defined but does not has a default prompt defined. Please do it like: '\"default\": \"prompt_with_instruction\"' to handle cases when a matching prompt can't be found.")
-            default_prompt_name = self.template.get("default")
-            if default_prompt_name not in self.template:
-                raise ValueError(
-                    f"The template {self.template_name} has \"default\" set to \"{default_prompt_name}\" but it's not defined. Please do it like: '\"{default_prompt_name}\": \"...\".")
-            prompt_name = get_prompt_name(variables, variable_names)
-            prompt_template = self.template.get(default_prompt_name)
-            if prompt_name in self.template:
-                prompt_template = self.template.get(prompt_name)
+            if self.template_module:
+                if type(variables) == list:
+                    variables = {k: v for k, v in zip(
+                        variable_names, variables)}
 
-            res = prompt_template.format(
-                **variables_to_dict(variables, variable_names))
+                res = self.template_module.get_prompt(variables)
+            else:
+                if type(variables) == dict:
+                    variables = [variables.get(name, None)
+                                 for name in variable_names]
+
+                if "default" not in self.template:
+                    raise ValueError(
+                        f"The template {self.template_name} has \"variables\" defined but does not has a default prompt defined. Please do it like: '\"default\": \"prompt_with_instruction\"' to handle cases when a matching prompt can't be found.")
+                default_prompt_name = self.template.get("default")
+                if default_prompt_name not in self.template:
+                    raise ValueError(
+                        f"The template {self.template_name} has \"default\" set to \"{default_prompt_name}\" but it's not defined. Please do it like: '\"{default_prompt_name}\": \"...\".")
+                prompt_name = get_prompt_name(variables, variable_names)
+                prompt_template = self.template.get(default_prompt_name)
+                if prompt_name in self.template:
+                    prompt_template = self.template.get(prompt_name)
+
+                res = prompt_template.format(
+                    **variables_to_dict(variables, variable_names))
 
         else:
             if type(variables) == dict:
@@ -92,18 +131,50 @@ class Prompter(object):
     def get_response(self, output: str) -> str:
         if self.template_name == "None":
             return output
+
+        splitted_output = output.split(self.template["response_split"])
+        # if len(splitted_output) <= 1:
+        #     return output.strip()
+
         return self.template["response_split"].join(
-            output.split(self.template["response_split"])[1:]
+            splitted_output[1:]
         ).strip()
 
     def get_variable_names(self) -> List[str]:
         if self.template_name == "None":
             return ["prompt"]
         elif "variables" in self.template:
-            return self.template.get("variables")
+            return self.template['variables']
         else:
             return ["instruction", "input"]
 
+    def get_train_data_from_dataset(self, data, only_first_n_items=None):
+        if self.template_module:
+            if hasattr(self.template_module,
+                       "get_train_data_list_from_dataset"):
+                data = self.template_module.get_train_data_list_from_dataset(
+                    data)
+            if only_first_n_items:
+                data = data[:only_first_n_items]
+            return list(itertools.chain(*list(
+                map(self.template_module.get_train_data, data)
+            )))
+
+        if only_first_n_items:
+            data = data[:only_first_n_items]
+
+        data = process_json_dataset(data)
+
+        train_data = [
+            {
+                'prompt': self.generate_prompt(d['variables']),
+                'completion': d['output'],
+                **{"_var_" + k: v for k, v in d['variables'].items()}
+            }
+            for d in data]
+
+        return train_data
+
 
 def get_val(arr, index, default=None):
     return arr[index] if -len(arr) <= index < len(arr) else default
@@ -116,4 +187,62 @@ def get_prompt_name(variables, variable_names):
 
 
 def variables_to_dict(variables, variable_names):
-    return {key: (variables[i] if i < len(variables) and variables[i] is not None else '') for i, key in enumerate(variable_names)}
+    return {
+        key: (variables[i] if i < len(variables)
+              and variables[i] is not None else '')
+        for i, key in enumerate(variable_names)
+    }
+
+
+def process_json_dataset(data):
+    if not isinstance(data, list):
+        raise ValueError("The dataset is not an array of objects.")
+
+    first_item = get_val_from_arr(data, 0, None)
+
+    if first_item is None:
+        raise ValueError("The dataset is empty.")
+    if not isinstance(first_item, dict):
+        raise ValueError("The dataset is not an array of objects.")
+
+    # Convert OpenAI fine-tuning dataset to LLaMA LoRA style
+    if "completion" in first_item and "output" not in first_item:
+        data = [
+            {"output" if k == "completion" else k: v for k, v in d.items()}
+            for d in data]
+        first_item = get_val_from_arr(data, 0, None)
+
+    # Flatten Stanford Alpaca style instances
+    if "instances" in first_item and isinstance(first_item["instances"], list):
+        data = [
+            {"output" if k == "completion" else k: v for k, v in d.items()}
+            for d in data]
+        flattened_data = []
+        for item in data:
+            for instance in item["instances"]:
+                d = {k: v for k, v in item.items() if k != "instances"}
+                d.update(instance)
+                flattened_data.append(d)
+        data = flattened_data
+        first_item = get_val_from_arr(data, 0, None)
+
+    if "output" not in first_item:
+        raise ValueError(
+            "The data does not contains an \"output\" or \"completion\".")
+
+    # Put all variables under the "variables" key if it does not exists
+    if "variables" not in first_item:
+        data = [
+            {
+                "variables":
+                    {k: v for k, v in d.items() if k != "output"},
+                "output":
+                    d["output"]
+            }
+            for d in data
+        ]
+    return data
+
+
+def get_val_from_arr(arr, index, default=None):
+    return arr[index] if -len(arr) <= index < len(arr) else default

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

@@ -0,0 +1,19 @@
+{
+  "num_train_epochs": 10,
+  "learning_rate": 0.0003,
+  "cutoff_len": 512,
+  "lora_r": 16,
+  "lora_alpha": 16,
+  "lora_dropout": 0.05,
+  "lora_target_modules": [
+    "q_proj",
+    "v_proj",
+    "k_proj",
+    "o_proj"
+  ],
+  "train_on_inputs": true,
+  "group_by_length": false,
+  "save_steps": 2000,
+  "save_total_limit": 10,
+  "logging_steps": 10
+}

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"
+}

lora_models/unhelpful-ai-v01/finetune_params.json

@@ -0,0 +1,19 @@
+{
+  "num_train_epochs": 16,
+  "learning_rate": 0.0003,
+  "cutoff_len": 512,
+  "lora_r": 12,
+  "lora_alpha": 32,
+  "lora_dropout": 0.05,
+  "lora_target_modules": [
+    "q_proj",
+    "v_proj",
+    "k_proj",
+    "o_proj"
+  ],
+  "train_on_inputs": false,
+  "group_by_length": false,
+  "save_steps": 500,
+  "save_total_limit": 5,
+  "logging_steps": 10
+}