diff --git a/src/lmflow/args.py b/src/lmflow/args.py
index f8ebf737d..5d87919a1 100644
--- a/src/lmflow/args.py
+++ b/src/lmflow/args.py
@@ -392,10 +392,13 @@ class DatasetArguments:
a boolean indicating whether to train on prompt for conversation datasets such as ShareGPT.
disable_conversation_bos_token: bool
- a boolean indicating whether to disable the bos token for conversation datasets.
+ [DEPRECATE SOON] a boolean indicating whether to disable the bos token for conversation datasets.
disable_conversation_eos_token: bool
- a boolean indicating whether to disable the eos token for conversation datasets.
+ [DEPRECATE SOON] a boolean indicating whether to disable the eos token for conversation datasets.
+
+ conversation_template: str
+ a string representing the template for conversation datasets.
The class also includes some additional parameters that can be used to configure the dataset further, such as `overwrite_cache`,
`validation_split_percentage`, `preprocessing_num_workers`, `disable_group_texts`, `demo_example_in_prompt`, `explanation_in_prompt`,
@@ -517,6 +520,10 @@ class DatasetArguments:
default=False,
metadata={"help": "Whether to disable the eos token for conversation datasets."}
)
+ conversation_template: Optional[str] = field(
+ default=None,
+ metadata={"help": "The template for conversation datasets."}
+ )
def __post_init__(self):
if self.streaming:
diff --git a/src/lmflow/datasets/dataset.py b/src/lmflow/datasets/dataset.py
index 343e6f416..36da7ce07 100644
--- a/src/lmflow/datasets/dataset.py
+++ b/src/lmflow/datasets/dataset.py
@@ -147,10 +147,17 @@ def _check_data_format(self):
correct_fields = INSTANCE_FIELDS_MAP[data_type]
# TODO: this can not guarantee every instance has correct fields.
if set(fields) != set(correct_fields):
- raise ValueError(
- f'Data instance fields incorrect'
- f' {list(fields)}: should be {list(correct_fields)}.'
- )
+ if data_type == "conversation":
+ if "messages" not in fields:
+ raise ValueError(
+ f'Conversation dataset should have "messages" field'
+ f' but got {list(fields)}'
+ )
+ else:
+ raise ValueError(
+ f'Data instance fields incorrect'
+ f' {list(fields)}: should be {list(correct_fields)}.'
+ )
def from_dict(self, dict_obj: dict, *args, **kwargs):
@@ -215,12 +222,19 @@ def from_dict(self, dict_obj: dict, *args, **kwargs):
for i, instance in enumerate(dict_obj[KEY_INSTANCES]):
fields = instance.keys()
if set(fields) != set(correct_fields):
- raise ValueError(
- f'data instance fields incorrect'
- f' {list(fields)}: should be {list(correct_fields)}.\n'
- f'The bad instance triggers the error, the {i}-th instance:\n'
- f' {instance}'
- )
+ if self.type == "conversation":
+ if "messages" not in fields:
+ raise ValueError(
+ f'Conversation dataset should have "messages" field'
+ f' but got {list(fields)}'
+ )
+ else:
+ raise ValueError(
+ f'data instance fields incorrect'
+ f' {list(fields)}: should be {list(correct_fields)}.\n'
+ f'The bad instance triggers the error, the {i}-th instance:\n'
+ f' {instance}'
+ )
try:
hf_dict = {}
diff --git a/src/lmflow/models/hf_decoder_model.py b/src/lmflow/models/hf_decoder_model.py
index 99d626bfd..68011110c 100644
--- a/src/lmflow/models/hf_decoder_model.py
+++ b/src/lmflow/models/hf_decoder_model.py
@@ -57,8 +57,15 @@
from lmflow.utils.constants import (
TEXT_ONLY_DATASET_DESCRIPTION,
TEXT2TEXT_DATASET_DESCRIPTION,
+ CONVERSATION_DATASET_DESCRIPTION
+)
+from lmflow.utils.conversation_template import (
+ ConversationTemplate,
+ EmptyConversationTemplate,
+ Llama2ConversationTemplate,
+ Qwen2ConversationTemplate,
+ EmptyConversationTemplateWithoutSpecialTokens
)
-from lmflow.utils.conversation_template import ConversationTemplate
from lmflow.utils.constants import CONVERSATION_ROLE_NAMES
@@ -420,6 +427,16 @@ def tokenize(self, dataset, add_special_tokens=True, *args, **kwargs):
)
dataset_type = dataset.get_type()
+ model_args = self.model_args
+ raw_datasets = dataset
+ hf_raw_datasets = dataset.get_backend_dataset()
+ column_names = list(hf_raw_datasets.features)
+
+ # since this will be pickled to avoid _LazyModule error in Hasher force
+ # logger loading before tokenize_function
+ tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
+
+ data_args = raw_datasets.get_data_args()
# Requires three types of information for tokenizing different datasets
# 1) Which fields require tokenization, e.g.
@@ -441,27 +458,40 @@ def tokenize(self, dataset, add_special_tokens=True, *args, **kwargs):
label_columns = ["output"]
add_special_tokens = False
elif dataset_type == "conversation":
- conversation_template: ConversationTemplate = kwargs.get("conversation_template", ConversationTemplate())
- logger.info(f"Conversation template: {conversation_template}")
+ conversation_template: ConversationTemplate = kwargs.get("conversation_template")
+ if conversation_template:
+ if data_args.conversation_template:
+ logger.warning("You specified conversation_template in both model.tokenize() and data_args. "
+ "Template in model.tokenize() will be used.")
+ else:
+ if data_args.conversation_template:
+ # TODO: make a registry for conversation templates
+ if data_args.conversation_template == 'llama2':
+ conversation_template = Llama2ConversationTemplate()
+ elif data_args.conversation_template == 'qwen2':
+ conversation_template = Qwen2ConversationTemplate()
+ elif data_args.conversation_template == 'empty':
+ conversation_template = EmptyConversationTemplate()
+ elif data_args.conversation_template == 'empty_no_special_tokens':
+ conversation_template = EmptyConversationTemplateWithoutSpecialTokens()
+ else:
+ raise NotImplementedError(
+ f"Conversation template {data_args.conversation_template} is not supported yet."
+ )
+ else:
+ logger.warning("No conversation template provided. Using default template.")
+ conversation_template = EmptyConversationTemplate()
+
+ logger.warning(f"Conversation template: {conversation_template}")
else:
raise NotImplementedError(
f"dataset type \"{dataset_type}\" is not supported, currently"
" only support following data types:\n"
f" 1) {TEXT_ONLY_DATASET_DESCRIPTION}\n"
f" 2) {TEXT2TEXT_DATASET_DESCRIPTION}\n"
+ f" 3) {CONVERSATION_DATASET_DESCRIPTION}\n"
)
- model_args = self.model_args
- raw_datasets = dataset
- hf_raw_datasets = dataset.get_backend_dataset()
- column_names = list(hf_raw_datasets.features)
-
- # since this will be pickled to avoid _LazyModule error in Hasher force
- # logger loading before tokenize_function
- tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
-
- data_args = raw_datasets.get_data_args()
-
# Whether to truncate long sequences to fit into max_length
use_truncation = False
if model_args.use_lora or data_args.disable_group_texts:
@@ -478,6 +508,8 @@ def tokenize_function(examples):
if dataset_type == "conversation":
for i in range(len(examples["messages"])):
messages = examples["messages"][i]
+ system = examples.get("system", [None] * num_example)[i]
+ tools = examples.get("tools", [None] * num_example)[i]
if len(messages) < 2 or messages[0]['role'] != CONVERSATION_ROLE_NAMES['user']:
tok_logger.warning(
"Invalid instance encountered. Either the conversation has less than "
@@ -491,13 +523,19 @@ def tokenize_function(examples):
)
messages = messages[:-1]
+ # DEPRECATION WARNING:
+ if data_args.disable_conversation_bos_token or data_args.disable_conversation_eos_token:
+ logger.warning(
+ "The disable_conversation_bos_token and disable_conversation_eos_token "
+ "flags are deprecated and will be removed in 2 versions. Please "
+ "customize your template and pass it through `conversation_template` "
+ "argument when calling .tokenize() instead.")
+
encoded_conversation = conversation_template.encode_conversation(
tokenizer=self.tokenizer,
messages=messages,
- system=examples["system"][i],
- tools=examples["tools"][i],
- disable_conversation_bos_token=data_args.disable_conversation_bos_token,
- disable_conversation_eos_token=data_args.disable_conversation_eos_token
+ system=system,
+ tools=tools,
)
input_ids, labels = [], []
diff --git a/src/lmflow/utils/constants.py b/src/lmflow/utils/constants.py
index 0cfb79965..616a4fbd8 100644
--- a/src/lmflow/utils/constants.py
+++ b/src/lmflow/utils/constants.py
@@ -78,6 +78,57 @@
).lstrip("\n")
+CONVERSATION_DATASET_DESCRIPTION = (
+"""
+"conversation": a dataset with conversation instances, with following format (`conversation_id`, `system` and `tools` are optional):
+
+ {
+ "type": "conversation",
+ "instances": [
+ {
+ "conversation_id": "CONVERSATION_ID",
+ "system": "SYSTEM_PROPMT",
+ "tools": ["TOOL_DESCRIPTION_1","TOOL_DESCRIPTION_2","TOOL_DESCRIPTION_X"],
+ "messages": [
+ {
+ "role": "user",
+ "content": "USER_INPUT_1"
+ },
+ {
+ "role": "assistant",
+ "content": "ASSISTANT_RESPONSE_1"
+ },
+ {
+ "role": "user",
+ "content": "USER_INPUT_2"
+ },
+ {
+ "role": "assistant",
+ "content": "ASSISTANT_RESPONSE_2"
+ }
+ ]
+ },
+ {
+ "conversation_id": "CONVERSATION_ID",
+ "system": "SYSTEM_PROPMT",
+ "tools": ["TOOL_DESCRIPTION_1"],
+ "messages": [
+ {
+ "role": "user",
+ "content": "USER_INPUT_1"
+ },
+ {
+ "role": "assistant",
+ "content": "ASSISTANT_RESPONSE_1"
+ }
+ ]
+ }
+ ]
+ }
+"""
+).lstrip("\n")
+
+
TEXT2TEXT_DATASET_DETAILS = (
"""
For example,
@@ -166,7 +217,7 @@
INSTANCE_FIELDS_MAP = {
"text_only": ["text"],
"text2text": ["input", "output"],
- "conversation": ["conversation_id", "system", "tools", "messages"],
+ "conversation": ["messages"], # system, tools and conversation_id are optional
"float_only": ["value"],
"image_text": ["images", "text"],
}
diff --git a/src/lmflow/utils/conversation_formatter.py b/src/lmflow/utils/conversation_formatter.py
index 313fad8f5..96d491977 100644
--- a/src/lmflow/utils/conversation_formatter.py
+++ b/src/lmflow/utils/conversation_formatter.py
@@ -1,2 +1,111 @@
-# reserved for handling different conversation formats,
-# such that users don't need to preprocess system prompt, tools, model-relate formats, etc.
\ No newline at end of file
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, Set, Sequence, Literal, Union, List, Optional
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TemplateComponent:
+ type: Literal['token', 'token_id', 'string', 'tools']
+ content: Union[str, int, List[str], List[int]]
+ mask: Optional[bool] = True # for token specific masking, work in progress
+
+ def __post_init__(self):
+ assert self.content, "Content of the component cannot be empty."
+
+ if self.type == 'tools':
+ assert isinstance(self.content, list), (
+ f"Content of tools component must be a list, got {type(self.content)}")
+ elif self.type in ['token', 'string']:
+ assert isinstance(self.content, str), (
+ f"Content of string/token component must be a string, got {type(self.content)}")
+ elif self.type == 'token_id':
+ assert isinstance(self.content, int) or all(isinstance(token_id, int) for token_id in self.content), (
+ f"Content of token_id component must be an integer or a list of integers.")
+ else:
+ raise ValueError(f"The type of the component must be either "
+ f"'token', 'string' or 'tools', got {self.type}")
+
+ def __repr__(self) -> str:
+ return f"TemplateComponent(type={self.type}, content={self.content})"
+
+ def __str__(self) -> str:
+ return f"{self.content}"
+
+
+@dataclass
+class Formatter(ABC):
+ template: List[TemplateComponent] = field(default_factory=list)
+
+ @abstractmethod
+ def format(self, **kwargs) -> List[TemplateComponent]: ...
+
+ def has_placeholder(self):
+ flag = False
+ for component in self.template:
+ if component.type == 'string':
+ if re.search(r"{{(.*?)}}", component.content):
+ flag = True
+ break
+ return flag
+
+
+@dataclass
+class EmptyFormatter(Formatter):
+ def __post_init__(self):
+ if self.has_placeholder():
+ raise ValueError("Empty formatter should not have placeholders.")
+
+ def format(self, **kwargs) -> list:
+ """Empty formatter for when no formatting is needed.
+ This is useful when user has already applied formatting to the dataset.
+
+ Returns
+ -------
+ list
+ Original template.
+ """
+ return self.template
+
+
+@dataclass
+class StringFormatter(Formatter):
+ def __post_init__(self):
+ if not self.has_placeholder():
+ raise ValueError("String formatter should have placeholders.")
+
+ def format(self, **kwargs) -> list:
+ """Format the string components with the provided keyword arguments.
+ Mostly used for formatting system prompt, user and assistant messages.
+
+ Parameters
+ ----------
+ **kwargs : dict
+ Keyword arguments containing values to replace in the template components.
+
+ Returns
+ -------
+ list
+ Formatted template.
+ """
+ formatted_template = []
+ for component in self.template:
+ if component.type == 'string':
+ for key, value in kwargs.items():
+ templated = component.content.replace("{{" + key + "}}", value)
+ formatted_template.append(TemplateComponent(type='string', content=templated))
+ else:
+ formatted_template.append(component)
+
+ logger.debug(formatted_template)
+ return formatted_template
+
+
+@dataclass
+class ListFormatter(Formatter):
+ def format(self, **kwargs) -> list:
+ pass # Work in progress
\ No newline at end of file
diff --git a/src/lmflow/utils/conversation_template.py b/src/lmflow/utils/conversation_template.py
index 9068eb717..43fb4f501 100644
--- a/src/lmflow/utils/conversation_template.py
+++ b/src/lmflow/utils/conversation_template.py
@@ -1,27 +1,34 @@
import logging
from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Dict, List, Optional, Sequence, Tuple, Union
from transformers import PreTrainedTokenizer
-from .constants import CONVERSATION_ROLE_NAMES
+from .conversation_formatter import Formatter, TemplateComponent, StringFormatter, EmptyFormatter
+
logger = logging.getLogger(__name__)
@dataclass
class ConversationTemplate:
+ user_formatter: Formatter
+ assistant_formatter: Formatter
+ system_formatter: Optional[Formatter] = None
+ tools_formatter: Optional[Formatter] = None
+ separator: Optional[TemplateComponent] = None
+
def encode_conversation(
self,
tokenizer: PreTrainedTokenizer,
messages: List[Dict[str, str]],
system: Optional[str] = None,
tools: Optional[List[str]] = None,
+ remove_last_sep: bool = False,
**kwargs
) -> Sequence[Tuple[List[int], List[int]]]:
r'''
Messages here should be guaranteed to be in pairs, with the first message being the user message and the second message being the system message.
- TODO: Support for different models.
Data example:
```json
{
@@ -30,22 +37,42 @@ def encode_conversation(
"tools": ["tool_1_desc"],
"messages": [
{
- "role": "human",
- "content": "I have a M via API so"
+ "role": "user",
+ "content": "hi"
},
{
- "role": "gpt",
- "content": "To"
+ "role": "assistant",
+ "content": "Hello!"
}
]
}
```
- '''
- encoded_pairs = self.__encode(tokenizer, messages, system, tools, **kwargs)
+ '''
+ assert isinstance(messages, list), "Messages must be a list."
+
+ if tools:
+ logger.warning("Tools are not supported yet. Please include tools in the system message manually.")
+
+ if system:
+ if system.replace(" ",""):
+ if not self.system_formatter:
+ raise ValueError("Your dataset contains system message but no system formatter is provided. "
+ "Consider either providing a system formatter or removing system prompt from your dataset.")
+ else:
+ system = None
+
+ encoded_pairs = self._encode(tokenizer, messages, system, tools, **kwargs)
+
+ if self.separator and remove_last_sep:
+ # For models that require a separator between messages,
+ # user can include the seperator at the end of each template
+ # and specify the separator. Auto formatting will remove the
+ # last separator once user specifies this option.
+ encoded_pairs = self.remove_last_separator(encoded_pairs, tokenizer)
return encoded_pairs
- def __encode(
+ def _encode(
self,
tokenizer: PreTrainedTokenizer,
messages: List[Dict[str, str]],
@@ -55,31 +82,145 @@ def __encode(
) -> Sequence[Tuple[List[int], List[int]]]:
# TODO: truncation according to model max length
# TODO: make sure the last few tokens are "learnable", not masked with token_id = -100.
- bos = [] if kwargs.get("disable_conversation_bos_token", False) else [tokenizer.bos_token_id]
- eos = [] if kwargs.get("disable_conversation_eos_token", False) else [tokenizer.eos_token_id]
+
res_all = []
+ system_formatted = self.system_formatter.format(content=system) if system else []
+ system_encoded = self._encode_template(system_formatted, tokenizer)
+
for i in range(0, len(messages), 2):
user_message = messages[i]
- system_message = messages[i + 1]
+ assistant_message = messages[i + 1]
- user_input = user_message["content"]
- system_input = system_message["content"]
+ user_formatted = self.user_formatter.format(content=user_message["content"])
+ assistant_formatted = self.assistant_formatter.format(content=assistant_message["content"])
- user_encoded = tokenizer.encode(user_input, add_special_tokens=False)
- system_encoded = tokenizer.encode(system_input, add_special_tokens=False)
+ user_encoded = self._encode_template(user_formatted, tokenizer)
+ assistant_encoded = self._encode_template(assistant_formatted, tokenizer)
res_all.append((
- bos + user_encoded,
- system_encoded + eos
+ system_encoded + user_encoded if i == 0 else user_encoded,
+ assistant_encoded
))
return res_all
+
+ def _encode_template(
+ self,
+ template: List[TemplateComponent],
+ tokenizer: PreTrainedTokenizer,
+ **kwargs
+ ) -> List[int]:
+ """Encode template components into token ids.
+
+ Parameters
+ ----------
+ template : List[TemplateComponent]
+ Formatted template components.
+ tokenizer : PreTrainedTokenizer
+ Tokenizer to convert tokens into token ids.
+
+ Returns
+ -------
+ List[int]
+ Encoded token ids.
+ """
+ encoded_ids = []
+ for component in template:
+ if component.type == 'string':
+ if len(component.content) == 0:
+ logger.warning("Empty string component found in the template.")
+ continue
+ else:
+ encoded_ids += tokenizer.encode(component.content, add_special_tokens=False)
+ elif component.type == 'token':
+ if component.content == 'bos_token':
+ encoded_ids += [tokenizer.bos_token_id]
+ elif component.content == 'eos_token':
+ encoded_ids += [tokenizer.eos_token_id]
+ else:
+ encoded_ids += [tokenizer.convert_tokens_to_ids(component.content)]
+ elif component.type == 'token_id':
+ encoded_ids += [component.content] if isinstance(component.content, int) else component.content
+ else:
+ raise NotImplementedError(f"Component type {component.type} is not supported yet.")
+ return encoded_ids
+
+ def remove_last_separator(
+ self,
+ encoded_pairs: Sequence[Tuple[List[int], List[int]]],
+ tokenizer: PreTrainedTokenizer
+ ) -> Sequence[Tuple[List[int], List[int]]]:
+ last_assistant_msg = encoded_pairs[-1][1]
+ if self.separator.type == 'string':
+ separator_ids = tokenizer.encode(self.separator.content, add_special_tokens=False)
+ elif self.separator.type == 'token':
+ separator_ids = [tokenizer.convert_tokens_to_ids(self.separator.content)]
+ else:
+ raise NotImplementedError(f"Component type {self.separator.type} cannot be used as a separator.")
+
+ if len(separator_ids) > len(last_assistant_msg):
+ raise ValueError("Separator is longer than the last assistant message, please check.")
+
+ if last_assistant_msg[-len(separator_ids):] == separator_ids:
+ last_assistant_msg = last_assistant_msg[:-len(separator_ids)]
+
+ encoded_pairs[-1] = (encoded_pairs[-1][0], last_assistant_msg)
+
+ return encoded_pairs
+
+@dataclass
+class EmptyConversationTemplate(ConversationTemplate):
+ user_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='token', content='bos_token'),
+ TemplateComponent(type='string', content='{{content}}')
+ ]
+ )
+ assistant_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='{{content}}'),
+ TemplateComponent(type='token', content='eos_token')
+ ]
+ )
+
+
+@dataclass
+class EmptyConversationTemplateWithoutSpecialTokens(ConversationTemplate):
+ user_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='{{content}}')
+ ]
+ )
+ assistant_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='{{content}}')
+ ]
+ )
+
@dataclass
class Llama2ConversationTemplate(ConversationTemplate):
- def __encode(
+ user_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='token', content='bos_token'),
+ TemplateComponent(type='string', content='[INST] {{content}} [/INST]')
+ ]
+ )
+ assistant_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='{{content}}'),
+ TemplateComponent(type='token', content='eos_token')
+ ]
+ )
+ system_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='<>\n{{content}}\n<>\n\n')
+ ]
+ )
+
+ def _encode(
self,
tokenizer: PreTrainedTokenizer,
messages: List[Dict[str, str]],
@@ -87,6 +228,49 @@ def __encode(
tools: Optional[str] = None,
**kwargs
) -> Sequence[Tuple[List[int], List[int]]]:
- '''The system info is included in the first round of user input for Llama2.
- '''
- pass
\ No newline at end of file
+ if tools:
+ logger.warning("Formatted tools are not supported in Llama2, thus tools will be ignored. "
+ "If this is intended, please include tools in the system message manually.")
+
+ res_all = []
+
+ system_formatted = self.system_formatter.format(content=system) if system else []
+ system_formatted_text = "".join([component.content for component in system_formatted if component.type == 'string']) # HACK
+
+ for i in range(0, len(messages), 2):
+ user_message = messages[i]
+ assistant_message = messages[i + 1]
+
+ user_content = system_formatted_text + user_message["content"] if i == 0 else user_message["content"]
+ user_formatted = self.user_formatter.format(content=user_content)
+ assistant_formatted = self.assistant_formatter.format(content=assistant_message["content"])
+
+ user_encoded = self._encode_template(user_formatted, tokenizer)
+ assistant_encoded = self._encode_template(assistant_formatted, tokenizer)
+
+ res_all.append((
+ user_encoded,
+ assistant_encoded
+ ))
+
+ return res_all
+
+
+@dataclass
+class Qwen2ConversationTemplate(ConversationTemplate):
+ user_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='<|im_start|>user\n{{content}}<|im_end|>\n')
+ ]
+ )
+ assistant_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='<|im_start|>assistant\n{{content}}<|im_end|>\n')
+ ]
+ )
+ system_formatter: Formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='string', content='<|im_start|>system\n{{content}}<|im_end|>\n')
+ ]
+ )
+ separator: TemplateComponent = TemplateComponent(type='string', content='\n')
\ No newline at end of file
diff --git a/tests/models/test_hf_decoder_model.py b/tests/models/test_hf_decoder_model.py
index 350729d93..4f3411d10 100644
--- a/tests/models/test_hf_decoder_model.py
+++ b/tests/models/test_hf_decoder_model.py
@@ -27,6 +27,11 @@
TEXT_ONLY_DATASET_DESCRIPTION,
TEXT2TEXT_DATASET_DESCRIPTION,
)
+from lmflow.utils.conversation_template import (
+ EmptyConversationTemplate,
+ Llama2ConversationTemplate,
+ EmptyConversationTemplateWithoutSpecialTokens,
+)
SAMPLE_TEXT = "Defintion: In this task, we ask you to write an answer to a question that involves events that may be stationary (not changing over time) or transient (changing over time). For example, the sentence \"he was born in the U.S.\" contains a stationary event since it will last forever; however, \"he is hungry\" contains a transient event since it will remain true for a short period of time. Note that a lot of the questions could have more than one correct answer. We only need a single most-likely answer. Please try to keep your \"answer\" as simple as possible. Concise and simple \"answer\" is preferred over those complex and verbose ones. \\n Input: Sentence: It's hail crackled across the comm, and Tara spun to retake her seat at the helm. \nQuestion: Will the hail storm ever end? \\n Output: NA \\n\\n"
@@ -39,6 +44,96 @@
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1
]
+CONVERSATION_SINGLETURN = {
+ "system": "sysinfo",
+ "messages": [
+ {
+ "role": "user",
+ "content": "Hello"
+ },
+ {
+ "role": "assistant",
+ "content": "Hi!"
+ }
+ ]
+}
+
+CONVERSATION_SINGLETURN_LLAMA2 = {
+ "messages": [
+ {
+ "role": "user",
+ "content": "[INST] <>\nsysinfo\n<>\n\nHello [/INST]"
+ },
+ {
+ "role": "assistant",
+ "content": "Hi!"
+ }
+ ]
+}
+
+CONVERSATION_SINGLETURN_LLAMA2_IDS = [
+ (
+ [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 9675, 3888, 13,
+ 29966, 829, 14816, 29903, 6778, 13, 13, 10994, 518, 29914, 25580, 29962],
+ [6324, 29991, 2]
+ )
+]
+
+CONVERSATION_MULTITURN = {
+ "system": "sysinfo",
+ "messages": [
+ {
+ "role": "user",
+ "content": "Hello"
+ },
+ {
+ "role": "assistant",
+ "content": "Hi!"
+ },
+ {
+ "role": "user",
+ "content": "How are you?"
+ },
+ {
+ "role": "assistant",
+ "content": "I'm good, thanks!"
+ }
+ ]
+}
+
+CONVERSATION_MULTITURN_LLAMA2 = {
+ "messages": [
+ {
+ "role": "user",
+ "content": "[INST] <>\nsysinfo\n<>\n\nHello [/INST]"
+ },
+ {
+ "role": "assistant",
+ "content": "Hi!"
+ },
+ {
+ "role": "user",
+ "content": "[INST] How are you? [/INST]"
+ },
+ {
+ "role": "assistant",
+ "content": "I'm good, thanks!"
+ }
+ ]
+}
+
+CONVERSATION_MULTITURN_LLAMA2_IDS = [
+ (
+ [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 9675, 3888, 13,
+ 29966, 829, 14816, 29903, 6778, 13, 13, 10994, 518, 29914, 25580, 29962],
+ [6324, 29991, 2]
+ ),
+ (
+ [1, 518, 25580, 29962, 1128, 526, 366, 29973, 518, 29914, 25580, 29962],
+ [306, 29915, 29885, 1781, 29892, 3969, 29991, 2]
+ )
+]
+
test_encode_input = "Question: Which of the following is not true for myelinated nerve fibers: (A) Impulse through myelinated fibers is slower than non-myelinated fibers (B) Membrane currents are generated at nodes of Ranvier (C) Saltatory conduction of impulses is seen (D) Local anesthesia is effective only when the nerve is not covered by myelin sheath."
test_encode_output = [24361, 25, 9022, 286, 262, 1708, 318, 407, 2081, 329, 616, 417, 3898, 16384, 26742, 25, 357, 32, 8, 9855, 9615, 832, 616, 417, 3898, 26742, 318, 13611, 621, 1729, 12, 1820, 417, 3898, 26742, 357, 33, 8, 4942, 1671, 1531, 28629, 389, 7560, 379, 13760, 286, 23075, 49663, 357, 34, 8, 13754, 2870, 369, 11124, 286, 37505, 318, 1775, 357, 35, 8, 10714, 49592, 318, 4050, 691, 618, 262, 16384, 318, 407, 5017, 416, 616, 27176, 673, 776, 13]
test_decode_input = [24361, 25, 9022, 286, 262, 1708, 318, 407, 2081, 329, 616, 417, 3898, 16384, 26742, 25, 357, 32, 8, 9855, 9615, 832, 616, 417, 3898, 26742, 318, 13611, 621, 1729, 12, 1820, 417, 3898, 26742, 357, 33, 8, 4942, 1671, 1531, 28629, 389, 7560, 379, 13760, 286, 23075, 49663, 357, 34, 8, 13754, 2870, 369, 11124, 286, 37505, 318, 1775, 357, 35, 8, 10714, 49592, 318, 4050, 691, 618, 262, 16384, 318, 407, 5017, 416, 616, 27176, 673, 776, 13]
@@ -46,6 +141,30 @@
test_inference_output = "The following is a list of the most common causes of myelinated nerve fibers."
+
+def make_gt_from_conversation_ids(conversation_ids):
+ res = {"input_ids": [], "attention_mask": [], "labels": []}
+ for turn_idx, turn_content in enumerate(conversation_ids):
+ user_content = turn_content[0]
+ assistant_content = turn_content[1]
+ res["input_ids"].extend(user_content)
+ res["input_ids"].extend(assistant_content)
+ res['attention_mask'].extend([1] * len(user_content) + [1] * len(assistant_content))
+ res['labels'].extend([-100] * len(user_content))
+ res['labels'].extend(assistant_content)
+ return res
+
+
+def make_gt_from_conversation_ids_batch(batched_conversation_ids):
+ res = {"input_ids": [], "attention_mask": [], "labels": []}
+ for conversation_ids in batched_conversation_ids:
+ this_res = make_gt_from_conversation_ids(conversation_ids)
+ res["input_ids"].append(this_res["input_ids"])
+ res["attention_mask"].append(this_res["attention_mask"])
+ res["labels"].append(this_res["labels"])
+ return res
+
+
class HFDecoderModelTest(unittest.TestCase):
def _test_tokenize(
@@ -53,8 +172,9 @@ def _test_tokenize(
model_name,
groundtruth_dataset,
groundtruth_tokenized_dataset,
+ **kwargs
):
- data_args = DatasetArguments(dataset_path=None)
+ data_args = DatasetArguments(dataset_path=None, disable_group_texts=False)
dataset = Dataset(data_args, backend="huggingface")
dataset = dataset.from_dict(groundtruth_dataset)
@@ -63,7 +183,7 @@ def _test_tokenize(
model_args = ModelArguments(model_name_or_path=model_name)
model = HFDecoderModel(model_args)
- tokenized_dataset = model.tokenize(dataset)
+ tokenized_dataset = model.tokenize(dataset, **kwargs)
self.assertEqual(
tokenized_dataset.get_backend_dataset().to_dict(),
@@ -141,9 +261,6 @@ def test_tokenize_conversation(self):
"type": "conversation",
"instances": [
{
- "conversation_id": 1,
- "system": "sysinfo_1",
- "tools": ["tool_1_desc"],
"messages": [
{
"role": "user",
@@ -166,7 +283,22 @@ def test_tokenize_conversation(self):
self._test_tokenize(
model_name="gpt2",
groundtruth_dataset=conversation_dataset,
- groundtruth_tokenized_dataset=conversation_tokenized_dataset
+ groundtruth_tokenized_dataset=conversation_tokenized_dataset,
+ conversation_template=EmptyConversationTemplateWithoutSpecialTokens()
+ )
+
+ self._test_tokenize(
+ model_name='meta-llama/Llama-2-7b-hf',
+ groundtruth_dataset={"type": "conversation", "instances": [CONVERSATION_SINGLETURN_LLAMA2]},
+ groundtruth_tokenized_dataset=make_gt_from_conversation_ids_batch([CONVERSATION_SINGLETURN_LLAMA2_IDS]),
+ conversation_template=EmptyConversationTemplate()
+ )
+
+ self._test_tokenize(
+ model_name='meta-llama/Llama-2-7b-hf',
+ groundtruth_dataset={"type": "conversation", "instances": [CONVERSATION_SINGLETURN]},
+ groundtruth_tokenized_dataset=make_gt_from_conversation_ids_batch([CONVERSATION_SINGLETURN_LLAMA2_IDS]),
+ conversation_template=Llama2ConversationTemplate()
)
@@ -175,9 +307,6 @@ def test_tokenize_conversation_multiple(self):
"type": "conversation",
"instances": [
{
- "conversation_id": 1,
- "system": "sysinfo_1",
- "tools": ["tool_1_desc_1"],
"messages": [
{
"role": "user",
@@ -190,9 +319,6 @@ def test_tokenize_conversation_multiple(self):
]
},
{
- "conversation_id": 2,
- "system": "sysinfo_2",
- "tools": ["tool_2_desc_1"],
"messages": [
{
"role": "user",
@@ -215,7 +341,22 @@ def test_tokenize_conversation_multiple(self):
self._test_tokenize(
model_name="gpt2",
groundtruth_dataset=conversation_dataset,
- groundtruth_tokenized_dataset=conversation_tokenized_dataset
+ groundtruth_tokenized_dataset=conversation_tokenized_dataset,
+ conversation_template=EmptyConversationTemplateWithoutSpecialTokens()
+ )
+
+ self._test_tokenize(
+ model_name='meta-llama/Llama-2-7b-hf',
+ groundtruth_dataset={"type": "conversation", "instances": [CONVERSATION_MULTITURN_LLAMA2]},
+ groundtruth_tokenized_dataset=make_gt_from_conversation_ids_batch([CONVERSATION_MULTITURN_LLAMA2_IDS]),
+ conversation_template=EmptyConversationTemplate()
+ )
+
+ self._test_tokenize(
+ model_name='meta-llama/Llama-2-7b-hf',
+ groundtruth_dataset={"type": "conversation", "instances": [CONVERSATION_MULTITURN]},
+ groundtruth_tokenized_dataset=make_gt_from_conversation_ids_batch([CONVERSATION_MULTITURN_LLAMA2_IDS]),
+ conversation_template=Llama2ConversationTemplate()
)
@@ -241,7 +382,6 @@ def test_decode(self):
self.assertEqual(model.decode(batch_decode_input), batch_decode_output)
- # @unittest.skip("deepspeed master_port conflict")
def test_inference(self):
ds_config_path = "examples/ds_config.json"
with open (ds_config_path, "r") as f:
diff --git a/tests/utils/test_conversation_formatter.py b/tests/utils/test_conversation_formatter.py
new file mode 100644
index 000000000..6b94b0417
--- /dev/null
+++ b/tests/utils/test_conversation_formatter.py
@@ -0,0 +1,25 @@
+import unittest
+from lmflow.utils.conversation_formatter import StringFormatter, TemplateComponent
+
+
+class StringFormatterTest(unittest.TestCase):
+
+ def test_format_string_component(self):
+ formatter = StringFormatter(
+ template=[
+ TemplateComponent(type='token', content='bos_token'),
+ TemplateComponent(type='string', content='[INST] {{content}} [/INST]'),
+ TemplateComponent(type='token', content='eos_token')
+ ]
+ )
+ formatted_components = formatter.format(content='Who are you?')
+ expected_components = [
+ TemplateComponent(type='token', content='bos_token'),
+ TemplateComponent(type='string', content='[INST] Who are you? [/INST]'),
+ TemplateComponent(type='token', content='eos_token')
+ ]
+ self.assertEqual(formatted_components, expected_components)
+
+
+if __name__ == '__main__':
+ unittest.main()
\ No newline at end of file
diff --git a/tests/utils/test_conversation_template.py b/tests/utils/test_conversation_template.py
index d2f6384b6..798e2e7d9 100644
--- a/tests/utils/test_conversation_template.py
+++ b/tests/utils/test_conversation_template.py
@@ -3,12 +3,24 @@
from transformers import AutoTokenizer
-from src.lmflow.utils.conversation_template import ConversationTemplate
+from lmflow.utils.conversation_template import EmptyConversationTemplate, Llama2ConversationTemplate, Qwen2ConversationTemplate
-MODEL_PATH = 'meta-llama/Llama-2-7b-hf'
+CONVERSATION_SINGLETURN = {
+ "system": "sysinfo",
+ "messages": [
+ {
+ "role": "user",
+ "content": "Hello"
+ },
+ {
+ "role": "assistant",
+ "content": "Hi!"
+ }
+ ]
+}
-conversation_singleturn_llama2 = [
+CONVERSATION_SINGLETURN_LLAMA2 = [
{
"role": "user",
"content": "[INST] <>\nsysinfo\n<>\n\nHello [/INST]"
@@ -19,15 +31,37 @@
}
]
-conversation_singleturn_llama2_ids = [
+CONVERSATION_SINGLETURN_LLAMA2_IDS = [
(
- [0, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 9675, 3888, 13,
+ [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 9675, 3888, 13,
29966, 829, 14816, 29903, 6778, 13, 13, 10994, 518, 29914, 25580, 29962],
- [6324, 29991, 0]
+ [6324, 29991, 2]
)
]
-conversation_multiturn_llama2 = [
+CONVERSATION_MULTITURN = {
+ "system": "sysinfo",
+ "messages": [
+ {
+ "role": "user",
+ "content": "Hello"
+ },
+ {
+ "role": "assistant",
+ "content": "Hi!"
+ },
+ {
+ "role": "user",
+ "content": "How are you?"
+ },
+ {
+ "role": "assistant",
+ "content": "I'm good, thanks!"
+ }
+ ]
+}
+
+CONVERSATION_MULTITURN_LLAMA2 = [
{
"role": "user",
"content": "[INST] <>\nsysinfo\n<>\n\nHello [/INST]"
@@ -46,38 +80,107 @@
}
]
-conversation_multiturn_llama2_ids = [
+CONVERSATION_MULTITURN_LLAMA2_IDS = [
(
- [0, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 9675, 3888, 13,
+ [1, 518, 25580, 29962, 3532, 14816, 29903, 6778, 13, 9675, 3888, 13,
29966, 829, 14816, 29903, 6778, 13, 13, 10994, 518, 29914, 25580, 29962],
- [6324, 29991, 0]
+ [6324, 29991, 2]
),
(
- [0, 518, 25580, 29962, 1128, 526, 366, 29973, 518, 29914, 25580, 29962],
- [306, 29915, 29885, 1781, 29892, 3969, 29991, 0]
+ [1, 518, 25580, 29962, 1128, 526, 366, 29973, 518, 29914, 25580, 29962],
+ [306, 29915, 29885, 1781, 29892, 3969, 29991, 2]
)
]
-class TestConversationTemplate(unittest.TestCase):
+class EmptyConversationTemplateTest(unittest.TestCase):
def setUp(self):
+ MODEL_PATH = 'meta-llama/Llama-2-7b-hf'
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
- self.conversation_template = ConversationTemplate()
+ self.conversation_template = EmptyConversationTemplate()
def test_encode_conversation_singleturn_llama2(self):
res = self.conversation_template.encode_conversation(
tokenizer=self.tokenizer,
- messages=conversation_singleturn_llama2,
+ messages=CONVERSATION_SINGLETURN_LLAMA2,
system=None,
tools=None
)
- self.assertEqual(res, conversation_singleturn_llama2_ids)
+ self.assertEqual(res, CONVERSATION_SINGLETURN_LLAMA2_IDS)
def test_encode_conversation_multiturn_llama2(self):
res = self.conversation_template.encode_conversation(
tokenizer=self.tokenizer,
- messages=conversation_multiturn_llama2,
+ messages=CONVERSATION_MULTITURN_LLAMA2,
system=None,
tools=None
)
- self.assertEqual(res, conversation_multiturn_llama2_ids)
\ No newline at end of file
+ self.assertEqual(res, CONVERSATION_MULTITURN_LLAMA2_IDS)
+
+
+class Llama2ConversationTemplateTest(unittest.TestCase):
+ def setUp(self):
+ MODEL_PATH = 'meta-llama/Llama-2-7b-hf'
+ self.tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
+ self.conversation_template = Llama2ConversationTemplate()
+
+ def test_encode_conversation_singleturn(self):
+ res = self.conversation_template.encode_conversation(
+ tokenizer=self.tokenizer,
+ messages=CONVERSATION_SINGLETURN['messages'],
+ system=CONVERSATION_SINGLETURN['system'],
+ tools=None
+ )
+ self.assertEqual(res, CONVERSATION_SINGLETURN_LLAMA2_IDS)
+
+ def test_encode_conversation_multiturn(self):
+ res = self.conversation_template.encode_conversation(
+ tokenizer=self.tokenizer,
+ messages=CONVERSATION_MULTITURN['messages'],
+ system=CONVERSATION_MULTITURN['system'],
+ tools=None
+ )
+ self.assertEqual(res, CONVERSATION_MULTITURN_LLAMA2_IDS)
+
+
+class Qwen2ConversationTemplateTest(unittest.TestCase):
+ def setUp(self):
+ MODEL_PATH = 'Qwen/Qwen1.5-0.5B-Chat'
+ self.tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
+ self.conversation_template = Qwen2ConversationTemplate()
+
+ def test_encode_conversation_singleturn(self):
+ res = self.conversation_template.encode_conversation(
+ tokenizer=self.tokenizer,
+ messages=CONVERSATION_SINGLETURN['messages'],
+ system=CONVERSATION_SINGLETURN['system'],
+ tools=None
+ )
+ print(res)
+
+ def test_encode_conversation_multiturn(self):
+ res = self.conversation_template.encode_conversation(
+ tokenizer=self.tokenizer,
+ messages=CONVERSATION_MULTITURN['messages'],
+ system=CONVERSATION_MULTITURN['system'],
+ tools=None
+ )
+ print(res)
+
+ print('===')
+ print(self.tokenizer.apply_chat_template(
+ CONVERSATION_MULTITURN['messages'],
+ tokenize=True,
+ add_generation_prompt=False
+ ))
+ print('===')
+ print(self.tokenizer.apply_chat_template(
+ CONVERSATION_MULTITURN['messages'],
+ tokenize=False,
+ add_generation_prompt=False
+ ))
+ print('===')
+
+
+if __name__ == '__main__':
+ unittest.main()
\ No newline at end of file