support .py prompt template

llama_lora/lib/finetune.py

@@ -162,6 +162,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:
@@ -221,7 +223,7 @@ 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,
@@ -260,6 +262,14 @@ def train(
         }
         json.dump(finetune_params, finetune_params_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
 
     old_state_dict = model.state_dict

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,20 +394,21 @@ 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)
 

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