update

llama_lora/lib/finetune.py

@@ -70,7 +70,13 @@ def train(
     wandb_tags: List[str] = [],
     wandb_watch: str = "false",  # options: false | gradients | all
     wandb_log_model: str = "true",  # options: false | true
+    status_message_callback: Any = None,
 ):
+    if status_message_callback:
+        cb_result = status_message_callback("Preparing training...")
+        if cb_result:
+            return
+
     if lora_modules_to_save is not None and len(lora_modules_to_save) <= 0:
         lora_modules_to_save = None
 
@@ -171,6 +177,16 @@ def train(
     if ddp:
         device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
 
+    if status_message_callback:
+        if isinstance(base_model, str):
+            cb_result = status_message_callback(f"Preparing model '{base_model}' for training...")
+            if cb_result:
+                return
+        else:
+            cb_result = status_message_callback("Preparing model for training...")
+            if cb_result:
+                return
+
     model = base_model
     if isinstance(model, str):
         model_name = model
@@ -216,51 +232,16 @@ def train(
     # )
     tokenizer.padding_side = "left"  # Allow batched inference
 
-    def tokenize(prompt, add_eos_token=True):
-        # there's probably a way to do this with the tokenizer settings
-        # but again, gotta move fast
-        result = tokenizer(
-            prompt,
-            truncation=True,
-            max_length=cutoff_len,
-            padding=False,
-            return_tensors=None,
-        )
-        if (
-            result["input_ids"][-1] != tokenizer.eos_token_id
-            and len(result["input_ids"]) < cutoff_len
-            and add_eos_token
-        ):
-            result["input_ids"].append(tokenizer.eos_token_id)
-            result["attention_mask"].append(1)
-
-        result["labels"] = result["input_ids"].copy()
-
-        return result
-
-    def generate_and_tokenize_prompt(data_point):
-        full_prompt = data_point["prompt"] + data_point["completion"]
-        tokenized_full_prompt = tokenize(full_prompt)
-        if not train_on_inputs:
-            user_prompt = data_point["prompt"]
-            tokenized_user_prompt = tokenize(user_prompt, add_eos_token=False)
-            user_prompt_len = len(tokenized_user_prompt["input_ids"])
-
-            tokenized_full_prompt["labels"] = [
-                -100
-            ] * user_prompt_len + tokenized_full_prompt["labels"][
-                user_prompt_len:
-            ]  # could be sped up, probably
-        return tokenized_full_prompt
-
-    # will fail anyway.
     try:
         model = prepare_model_for_int8_training(model)
     except Exception as e:
         print(
             f"Got error while running prepare_model_for_int8_training(model), maybe the model has already be prepared. Original error: {e}.")
 
-    # model = prepare_model_for_int8_training(model)
+    if status_message_callback:
+        cb_result = status_message_callback("Preparing PEFT model for training...")
+        if cb_result:
+            return
 
     lora_config_args = {
         'r': lora_r,
@@ -279,12 +260,6 @@ def train(
     if bf16:
         model = model.to(torch.bfloat16)
 
-    # If train_data is a list, convert it to datasets.Dataset
-    if isinstance(train_data, list):
-        with open(os.path.join(output_dir, "train_data_samples.json"), 'w') as file:
-            json.dump(list(train_data[:100]), file, indent=2)
-        train_data = Dataset.from_list(train_data)
-
     if resume_from_checkpoint:
         # Check the available weights and load them
         checkpoint_name = os.path.join(
@@ -320,6 +295,54 @@ def train(
         wandb.config.update({"model": {"all_params": all_params, "trainable_params": trainable_params,
                             "trainable%": 100 * trainable_params / all_params}})
 
+    if status_message_callback:
+        cb_result = status_message_callback("Preparing train data...")
+        if cb_result:
+            return
+
+    def tokenize(prompt, add_eos_token=True):
+        # there's probably a way to do this with the tokenizer settings
+        # but again, gotta move fast
+        result = tokenizer(
+            prompt,
+            truncation=True,
+            max_length=cutoff_len,
+            padding=False,
+            return_tensors=None,
+        )
+        if (
+            result["input_ids"][-1] != tokenizer.eos_token_id
+            and len(result["input_ids"]) < cutoff_len
+            and add_eos_token
+        ):
+            result["input_ids"].append(tokenizer.eos_token_id)
+            result["attention_mask"].append(1)
+
+        result["labels"] = result["input_ids"].copy()
+
+        return result
+
+    def generate_and_tokenize_prompt(data_point):
+        full_prompt = data_point["prompt"] + data_point["completion"]
+        tokenized_full_prompt = tokenize(full_prompt)
+        if not train_on_inputs:
+            user_prompt = data_point["prompt"]
+            tokenized_user_prompt = tokenize(user_prompt, add_eos_token=False)
+            user_prompt_len = len(tokenized_user_prompt["input_ids"])
+
+            tokenized_full_prompt["labels"] = [
+                -100
+            ] * user_prompt_len + tokenized_full_prompt["labels"][
+                user_prompt_len:
+            ]  # could be sped up, probably
+        return tokenized_full_prompt
+
+    # If train_data is a list, convert it to datasets.Dataset
+    if isinstance(train_data, list):
+        with open(os.path.join(output_dir, "train_data_samples.json"), 'w') as file:
+            json.dump(list(train_data[:100]), file, indent=2)
+        train_data = Dataset.from_list(train_data)
+
     if val_set_size > 0:
         train_val = train_data.train_test_split(
             test_size=val_set_size, shuffle=True, seed=42
@@ -339,6 +362,11 @@ def train(
         model.is_parallelizable = True
         model.model_parallel = True
 
+    if status_message_callback:
+        cb_result = status_message_callback("Train starting...")
+        if cb_result:
+            return
+
     # https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments
     training_args = {
         'output_dir': output_dir,

llama_lora/ui/finetune/finetune_ui.py

@@ -309,6 +309,7 @@ def handle_lora_modules_to_save_add(choices, new_module, selected_modules):
 
 def do_abort_training():
     Global.should_stop_training = True
+    Global.training_status_text = "Aborting..."
 
 
 def finetune_ui():

llama_lora/ui/finetune/training.py

@@ -22,6 +22,13 @@ from ..trainer_callback import (
 from .data_processing import get_data_from_input
 
 
+def status_message_callback(message):
+    if Global.should_stop_training:
+        return True
+
+    Global.training_status_text = message
+
+
 def do_train(
     # Dataset
     template,
@@ -254,6 +261,7 @@ def do_train(
                 train_output = Global.finetune_train_fn(
                     train_data=train_data,
                     callbacks=training_callbacks,
+                    status_message_callback=status_message_callback,
                     **finetune_args,
                 )
 

llama_lora/ui/trainer_callback.py

@@ -57,6 +57,9 @@ def update_training_states(
     Global.training_log_history = log_history
     Global.training_eta = Global.training_eta_predictor.predict_eta(current_step, total_steps)
 
+    if Global.should_stop_training:
+        return
+
     last_history = None
     last_loss = None
     if len(Global.training_log_history) > 0: