fix additional tokens - rk

rjojjr · Aug 21, 2024 · a0f862f · a0f862f
1 parent 63b9bbe
commit a0f862f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/main/base/llm_base_module.py b/src/main/base/llm_base_module.py
@@ -133,7 +133,7 @@ def merge_base(arguments: MergeArguments, tokenizer, base_model, bnb_config) ->
     model_dir = f'{arguments.output_dir}/merged-models/{arguments.new_model}'
     print(f"merging {arguments.base_model} with LoRA into {arguments.new_model}")
 
-    if arguments.use_agent_tokens:
+    if arguments.use_agent_tokens or arguments.additional_vocabulary_tokens is not None:
         model = AutoPeftModelForCausalLM.from_pretrained(lora_dir)
     else:
         model = PeftModel.from_pretrained(base_model, lora_dir, quantization_config=bnb_config)

diff --git a/src/main/main.py b/src/main/main.py
@@ -88,6 +88,7 @@ def main() -> None:
 
     print('')
     print(f'Is Chat Model: {args.is_chat_model}')
+    print(f'Using Additional Vocab Tokens: {args.additional_vocabulary_tokens}')
     print(f'Is LangChain Agent Model: {args.use_agent_tokens}')
 
     if args.fine_tune: