[2025.06.06] update tokenizer train code

This commit is contained in:
coffee
2025-06-06 16:06:36 +08:00
parent 628aa8c674
commit f9e0ae5825
@@ -11,9 +11,22 @@ from tokenizers import (
Tokenizer, Tokenizer,
) )
from transformers import PreTrainedTokenizerFast from transformers import PreTrainedTokenizerFast
import logging
from tqdm import tqdm
log_dir = "/mnt/hanluzhi/dolly/logs"
os.makedirs(log_dir, exist_ok=True)
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(message)s',
handlers=[
logging.StreamHandler(),
logging.FileHandler(os.path.join(log_dir, "training.log"))
]
)
def get_training_corpus(): def get_training_corpus():
for file_path in jsonl_files: for file_path in tqdm(jsonl_files, desc="processing files", unit="file"):
with open(file_path, "r", encoding="utf-8") as f: with open(file_path, "r", encoding="utf-8") as f:
for line in f: for line in f:
data = json.loads(line) data = json.loads(line)
@@ -70,7 +83,7 @@ trainer = trainers.BpeTrainer(
show_progress=True show_progress=True
) )
print("Training tokenizer...") logging.info("Training tokenizer...")
tokenizer.train_from_iterator( tokenizer.train_from_iterator(
get_training_corpus(), get_training_corpus(),
trainer=trainer, trainer=trainer,
@@ -89,12 +102,12 @@ wrapped_tokenizer = PreTrainedTokenizerFast(
errors="replace", errors="replace",
split_special_tokens=False, split_special_tokens=False,
) )
logging.info("train done ")
output_dir = Path("./llm_tokenizer") output_dir = Path("./llm_tokenizer")
output_dir.mkdir(exist_ok=True, parents=True) output_dir.mkdir(exist_ok=True, parents=True)
wrapped_tokenizer.save_pretrained(output_dir) wrapped_tokenizer.save_pretrained(output_dir)
print(f"Tokenizer saved to {output_dir}") logging.info(f"Tokenizer saved to {output_dir}")
QWEN2_CHAT_TEMPLATE = ( QWEN2_CHAT_TEMPLATE = (
"{% for message in messages %}" "{% for message in messages %}"
@@ -123,14 +136,14 @@ formatted = wrapped_tokenizer.apply_chat_template(
tokenize=False, tokenize=False,
add_generation_prompt=True add_generation_prompt=True
) )
print("\n格式化后的对话:") logging.info("\n格式化后的对话:")
print(formatted) logging.info(formatted)
test_text = "深度学习是人工智能的一个重要分支。Deep learning is a subset of machine learning." test_text = "深度学习是人工智能的一个重要分支。Deep learning is a subset of machine learning."
print("\nTest tokenization:") logging.info("\nTest tokenization:")
encoded = wrapped_tokenizer.tokenize(test_text) encoded = wrapped_tokenizer.tokenize(test_text)
print(encoded) logging.info(encoded)
print(encoded) logging.info(encoded)
decoded_text = wrapped_tokenizer.decode(encoded) decoded_text = wrapped_tokenizer.decode(encoded)
print(decoded_text) logging.info(decoded_text)
print(decoded_text) logging.info(decoded_text)