Merge branch 'dev-2' of github.com:zetavg/llama-lora into dev-2

.gitignore

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

README.md

@@ -60,13 +60,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:

app.py

@@ -5,21 +5,37 @@ 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 = "",
     # 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 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 (
@@ -34,12 +50,22 @@ def main(
     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()
 

llama_lora/globals.py

@@ -40,6 +40,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
@@ -50,8 +51,66 @@ def train(
     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,
+        '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 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 +197,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:
@@ -197,15 +258,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 +278,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
 

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,65 @@
+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
+    }
+
+    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["stopping_criteria"].insert(
+                0,
+                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:
+                decoded_output = tokenizer.decode(output, skip_special_tokens=True)
+                yield decoded_output, output
+                if output[-1] in [tokenizer.eos_token_id]:
+                    break
+        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=True)
+    yield decoded_output, output
+    return

llama_lora/models.py

@@ -8,19 +8,7 @@ from transformers import LlamaForCausalLM, 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):
@@ -60,9 +48,10 @@ def get_new_base_model(base_model_name):
             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)
+    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
 

llama_lora/ui/finetune_ui.py

@@ -79,56 +79,50 @@ 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
 
 
@@ -144,75 +138,59 @@ def refresh_preview(
     preview_show_actual_prompt,
 ):
     try:
-        max_preview_count = 100
+        max_preview_count = 30
         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
+        )
 
-        else:  # Load dataset from data directory
-            data = get_dataset_content(dataset_from_data_dir)
-            data = process_json_dataset(data)
+        train_data = prompter.get_train_data_from_dataset(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)]
+
 
-        headers = headers + ["Completion (output)"]
-        preview_data = result = [pd + [d['output']]
-                                 for pd, d in zip(preview_data, data[:max_preview_count])]
+        # 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)]
 
-        preview_info_message = f"The dataset has a total of {data_count} item(s)."
+        # headers = headers + ["Completion (output)"]
+        # preview_data = result = [pd + [d['output']]
+        #                          for pd, d in zip(preview_data, data[:max_preview_count])]
+
+        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"{data_count} item(s)."
         if load_dataset_from == "Data Dir":
-            info_message = "This dataset contains " + info_message
+            info_message = "This dataset contains about " + 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
@@ -288,57 +266,24 @@ 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
-                )
+        # 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
+        )
 
-        else:  # Load dataset from data directory
-            data = get_dataset_content(dataset_from_data_dir)
-            data = process_json_dataset(data)
+        train_data = prompter.get_train_data_from_dataset(data)
 
-        data_count = len(data)
+        data_count = len(train_data)
         evaluate_data_count = math.ceil(data_count * evaluate_data_percentage)
 
-        train_data = [
-            {
-                'prompt': prompter.generate_prompt(d['variables']),
-                'completion': d['output']}
-            for d in data]
-
         def get_progress_text(epoch, epochs, last_loss):
             progress_detail = f"Epoch {math.ceil(epoch)}/{epochs}"
             if last_loss is not None:
@@ -449,26 +394,33 @@ 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_percentage': evaluate_data_percentage,
 
-                '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,
             }
             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
@@ -491,7 +443,12 @@ Train data (first 10):
             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)

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()
 
@@ -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,26 +111,56 @@ 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
-        )
+
+        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
+        }
+
+        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)
+
+            if Global.should_stop_generating:
+                return
+
+            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
+
 
         if stream_output:
             # Stream the reply 1 token at a time.

llama_lora/utils/data.py

@@ -30,7 +30,7 @@ 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")]
+    return [filename.rstrip(".json") for filename in all_files if fnmatch.fnmatch(filename, "*.json") or fnmatch.fnmatch(filename, "*.py")]
 
 
 def get_available_dataset_names():

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:
@@ -104,6 +143,30 @@ class Prompter(object):
         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
@@ -117,3 +180,57 @@ 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)}
+
+
+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