fix baichuan template

Former-commit-id: 9c6dd1051417c91074daa7dd6ed6cc53448135ad
This commit is contained in:
hiyouga 2023-08-11 13:45:47 +08:00
parent abdfa26d06
commit f1485ab927
2 changed files with 12 additions and 7 deletions

View File

@ -1,3 +1,4 @@
import tiktoken
from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal
from itertools import chain
@ -31,10 +32,11 @@ def preprocess_dataset(
def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, Any]:
# build grouped texts with format `X1 X2 X3 ...` (without <eos>)
if hasattr(tokenizer, "tokenizer"): # for tiktoken tokenizer (Qwen)
if isinstance(getattr(tokenizer, "tokenizer"), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
kwargs = dict(allowed_special="all")
else:
kwargs = dict(add_special_tokens=False)
tokenized_examples = tokenizer(examples["prompt"], **kwargs)
concatenated_examples = {k: list(chain(*tokenized_examples[k])) for k in tokenized_examples.keys()}
total_length = len(concatenated_examples[list(concatenated_examples.keys())[0]])

View File

@ -1,5 +1,6 @@
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
import tiktoken
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
from llmtuner.extras.logging import get_logger
@ -122,7 +123,7 @@ class Template:
r"""
Converts context to token ids.
"""
if hasattr(tokenizer, "tokenizer"): # for tiktoken tokenizer (Qwen)
if isinstance(getattr(tokenizer, "tokenizer"), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
kwargs = dict(allowed_special="all")
else:
kwargs = dict(add_special_tokens=False)
@ -428,14 +429,16 @@ Supports: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
"""
register_template(
name="baichuan",
prefix=[],
prefix=[
{"token": "<reserved_102>"} # user token (a little difference in the first turn)
],
prompt=[
{"token": "<reserved_102>"}, # user token (a little difference in position)
"{{query}}"
"{{query}}",
{"token": "<reserved_103>"} # assistant token
],
sep=[],
stop_words=[
"<reserved_103>" # assistant token
{"token": "<reserved_102>"} # user token
],
use_history=True
)