save train info with model

llama_lora/lib/finetune.py

@@ -2,6 +2,8 @@ import os
 import sys
 from typing import Any, List
 
+import json
+
 import fire
 import torch
 import transformers
@@ -47,6 +49,10 @@ def train(
     # logging
     callbacks: List[Any] = []
 ):
+    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(f"The output directory already exists and is not empty. ({output_dir})")
+
     device_map = "auto"
     world_size = int(os.environ.get("WORLD_SIZE", 1))
     ddp = world_size != 1
@@ -202,6 +208,12 @@ def train(
         ),
         callbacks=callbacks,
     )
+
+    if not os.path.exists(output_dir):
+        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)
+
     model.config.use_cache = False
 
     old_state_dict = model.state_dict
@@ -214,9 +226,16 @@ def train(
     if torch.__version__ >= "2" and sys.platform != "win32":
         model = torch.compile(model)
 
-    result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    train_output = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
 
     model.save_pretrained(output_dir)
     print(f"Model saved to {output_dir}.")
 
-    return result
+    with open(os.path.join(output_dir, "trainer_log_history.jsonl"), 'w') as trainer_log_history_jsonl_file:
+        trainer_log_history = "\n".join([json.dumps(line) for line in trainer.state.log_history])
+        trainer_log_history_jsonl_file.write(trainer_log_history)
+
+    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)
+
+    return train_output

llama_lora/ui/finetune_ui.py

@@ -419,11 +419,27 @@ Train data (first 10):
         # Do not let other tqdm iterations interfere the progress reporting after training starts.
         # progress.track_tqdm = False  # setting this dynamically is not working, determining if track_tqdm should be enabled based on GPU cores at start instead.
 
+        output_dir = os.path.join(Global.data_dir, "lora_models", model_name)
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        with open(os.path.join(output_dir, "info.json"), 'w') as info_json_file:
+            dataset_name = "N/A (from text input)"
+            if load_dataset_from == "Data Dir":
+                dataset_name = dataset_from_data_dir
+
+            info = {
+                'base_model': Global.base_model,
+                'prompt_template': template,
+                'dataset_name': dataset_name,
+                'dataset_rows': len(train_data),
+            }
+            json.dump(info, info_json_file, indent=2)
+
         results = Global.train_fn(
             base_model,  # base_model
             tokenizer,  # tokenizer
-            os.path.join(Global.data_dir, "lora_models",
-                         model_name),  # output_dir
+            output_dir,  # output_dir
             train_data,
             # 128,  # batch_size (is not used, use gradient_accumulation_steps instead)
             micro_batch_size,    # micro_batch_size
@@ -451,7 +467,7 @@ Train data (first 10):
         return result_message
 
     except Exception as e:
-        raise gr.Error(e)
+        raise gr.Error(f"{e} (To dismiss this error, click the 'Abort' button)")
 
 
 def do_abort_training():