From f9793654ff7030180eda1e5da417938631212b6e Mon Sep 17 00:00:00 2001
From: Robert Kirchner <rjojjr91@live.com>
Date: Sun, 11 Aug 2024 11:50:43 -0500
Subject: [PATCH 1/2] fix agent tokens merge - rk

---
 src/main/base/llm_base_module.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/main/base/llm_base_module.py b/src/main/base/llm_base_module.py
index 32ed326..65a16ed 100644
--- a/src/main/base/llm_base_module.py
+++ b/src/main/base/llm_base_module.py
@@ -1,6 +1,6 @@
 from arguments.arguments import TuneArguments, MergeArguments, PushArguments
 from datasets import load_dataset
-from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training, TaskType, PeftConfig
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType, AutoPeftModelForCausalLM, PeftModel
 from trl import SFTTrainer, SFTConfig, setup_chat_format
 from transformers.trainer_utils import get_last_checkpoint
 from utils.model_utils import get_all_layers, get_all_linear_layers
@@ -12,11 +12,11 @@
 
 
 def _add_agent_tokens(tokenizer, model):
-    agent_tokens = ["Thought:", "Action:", "Input:", "Observation:", "Answer:", "Action\sInput:", "Final\sAnswer:"]
-
+    agent_tokens = ["Thought:", "Action:", "Action Input:", "Observation:", "Final Answer:"]
     agent_tokens = set(agent_tokens) - set(tokenizer.vocab.keys())
     tokenizer.add_tokens(list(agent_tokens))
-    model.resize_token_embeddings(len(tokenizer))
+    if model is not None:
+        model.resize_token_embeddings(len(tokenizer))
 
 
 # TODO - Tune/extract an embeddings only model
@@ -135,10 +135,12 @@ def merge_base(arguments: MergeArguments, tokenizer, base_model, bnb_config) ->
     lora_dir = f"{arguments.output_dir}/checkpoints/{arguments.new_model}/adapter"
     model_dir = f'{arguments.output_dir}/{arguments.new_model}'
     print(f"merging {arguments.base_model} with LoRA into {arguments.new_model}")
-    print('')
 
-    config = PeftConfig.from_pretrained(lora_dir)
-    model = PeftModel.from_pretrained(base_model, lora_dir, quantization_config=bnb_config, config=config)
+    if arguments.use_agent_tokens:
+        model = AutoPeftModelForCausalLM.from_pretrained(lora_dir)
+    else:
+        model = PeftModel.from_pretrained(base_model, lora_dir, quantization_config=bnb_config)
+
     model = model.merge_and_unload(progressbar=True)
     print('')
 

From 2b1cc687b36746722645875b8266a308f911bfac Mon Sep 17 00:00:00 2001
From: Robert Kirchner <rjojjr91@live.com>
Date: Sun, 11 Aug 2024 11:51:22 -0500
Subject: [PATCH 2/2] bump patch - rk

---
 src/main/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main/main.py b/src/main/main.py
index c7bc0e2..928402e 100644
--- a/src/main/main.py
+++ b/src/main/main.py
@@ -9,7 +9,7 @@
 import os
 
 # TODO - Automate this
-version = '1.4.2'
+version = '1.4.3'
 
 # TODO - Change this once support for more LLMs is added
 title = f'Llama AI LLM LoRA Torch Text Fine-Tuner v{version}'