Fix NPU LLM example save/load tokenizer

intel · Dec 3, 2024 · c5fef1a · c5fef1a
1 parent ab01753
commit c5fef1a
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 6 deletions.
diff --git a/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/baichuan2.py b/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/baichuan2.py
@@ -79,6 +79,8 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
             transpose_value_cache=not args.disable_transpose_value_cache,
             save_directory=args.save_directory
         )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer.save_pretrained(args.save_directory)
     else:
         model = AutoModelForCausalLM.load_low_bit(
             args.save_directory,
@@ -90,8 +92,8 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
             transpose_value_cache=not args.disable_transpose_value_cache,
             trust_remote_code=True,
         )
+        tokenizer = AutoTokenizer.from_pretrained(args.save_directory, trust_remote_code=True)        
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     DEFAULT_SYSTEM_PROMPT = """\
     """

diff --git a/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/generate.py b/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/generate.py
@@ -43,7 +43,6 @@
     args = parser.parse_args()
     model_path = args.repo_id_or_model_path
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     if not args.lowbit_path or not os.path.exists(args.lowbit_path):
         model = AutoModelForCausalLM.from_pretrained(
@@ -52,13 +51,16 @@
             load_in_low_bit=args.load_in_low_bit,
             attn_implementation="eager"
         )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer.save_pretrained(args.lowbit_path)
     else:
         model = AutoModelForCausalLM.load_low_bit(
             args.lowbit_path,
             trust_remote_code=True,
             bigdl_transformers_low_bit=args.load_in_low_bit,
             attn_implementation="eager"
         )
+        tokenizer = AutoTokenizer.from_pretrained(args.lowbit_path, trust_remote_code=True)        
 
     print(model)
 

diff --git a/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/llama2.py b/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/llama2.py
@@ -79,6 +79,8 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
             transpose_value_cache=not args.disable_transpose_value_cache,
             save_directory=args.save_directory
         )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer.save_pretrained(args.save_directory)
     else:
         model = AutoModelForCausalLM.load_low_bit(
             args.save_directory,
@@ -89,8 +91,8 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
             max_prompt_len=args.max_prompt_len,
             transpose_value_cache=not args.disable_transpose_value_cache,
         )
+        tokenizer = AutoTokenizer.from_pretrained(args.save_directory, trust_remote_code=True)        
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     DEFAULT_SYSTEM_PROMPT = """\
     """

diff --git a/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/llama3.py b/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/llama3.py
@@ -80,6 +80,8 @@ def get_prompt(user_input: str, chat_history: list[tuple[str, str]],
             transpose_value_cache=not args.disable_transpose_value_cache,
             save_directory=args.save_directory
         )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer.save_pretrained(args.save_directory)
     else:
         model = AutoModelForCausalLM.load_low_bit(
             args.save_directory,
@@ -90,8 +92,8 @@ def get_prompt(user_input: str, chat_history: list[tuple[str, str]],
             max_prompt_len=args.max_prompt_len,
             transpose_value_cache=not args.disable_transpose_value_cache,
         )
+        tokenizer = AutoTokenizer.from_pretrained(args.save_directory, trust_remote_code=True)        
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     DEFAULT_SYSTEM_PROMPT = """\
     """

diff --git a/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/minicpm.py b/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/minicpm.py
@@ -65,6 +65,8 @@
             transpose_value_cache=not args.disable_transpose_value_cache,
             save_directory=args.save_directory
         )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer.save_pretrained(args.save_directory)
     else:
         model = AutoModelForCausalLM.load_low_bit(
             args.save_directory,
@@ -76,7 +78,7 @@
             transpose_value_cache=not args.disable_transpose_value_cache,
             trust_remote_code=True,
         )
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(args.save_directory, trust_remote_code=True)        
 
     print("-" * 80)
     print("done")

diff --git a/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/qwen.py b/python/llm/example/NPU/HF-Transformers-AutoModels/LLM/qwen.py
@@ -71,6 +71,8 @@
             quantization_group_size=args.quantization_group_size,
             save_directory=args.save_directory
         )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer.save_pretrained(args.save_directory)
     else:
         model = AutoModelForCausalLM.load_low_bit(
             args.save_directory,
@@ -81,8 +83,8 @@
             max_prompt_len=args.max_prompt_len,
             transpose_value_cache=not args.disable_transpose_value_cache,
         )
+        tokenizer = AutoTokenizer.from_pretrained(args.save_directory, trust_remote_code=True)        
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     print("-" * 80)
     print("done")