add formatting for llama2chat

llmfoundry/data/finetuning/chat.py

@@ -0,0 +1,160 @@
+# Copyright 2022 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Dict, List
+
+
+class ChatFormatter:
+    """A class for formatting the chat history.
+
+    Args:
+        system: The system prompt. If None, a default ChatML-formatted prompt is used.
+        user: The user prompt. If None, a default ChatML value is used.
+        assistant: The assistant prompt. If None, a default ChatML value is used.
+        turn_joiner: The string to use to join turns. Defaults to a newline.
+
+    Attributes:
+        system: The system prompt.
+        user: The user prompt.
+        assistant: The assistant prompt.
+        response_prefix: The response prefix (anything before {} in the assistant format string)
+        response_suffix: The response suffix (anything after {} in the assistant format string)
+        turn_joiner: The string to use to join turns.
+    """
+
+    def __init__(self, system: str = None, user: str = None, assistant: str = None, turn_joiner: str = '\n') -> None:
+        self.system = system if system else '<|im_start|>system\nA conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.<|im_end|>\n'
+        self.user = user if user else '<|im_start|>user\n{}<|im_end|>\n'
+        self.assistant = assistant if assistant else '<|im_start|>assistant\n{}<|im_end|>\n'
+        self.response_prefix = self.assistant.split('{}')[0]
+        self.response_suffix = self.assistant.split('{}')[1] if len(self.assistant.split('{}')) > 1 else ''
+        self.turn_joiner = turn_joiner
+
+    def parse_structured_history(self, history: List[Dict[str, str]]) -> List[List[str]]:
+        """Parses a structured history into a list of lists of strings.
+
+        Currently expects Llama2 format, but easily extended to ShareGPT and other similar formats.
+
+        Args:
+            history: A list of dictionaries. Each dictionary contains two keys: 'role', (user, assistant, or system) and 'content', the message.
+
+        Returns:
+            The chat history as a list of lists of strings.
+
+        Mutates:
+            system: The system prompt (if a system prompt is present in the history)
+        """
+        if len(history) == 0:
+            return []
+        if history[0]['role'] == 'system':
+            if hasattr(self, "system_fmt"):
+                self.system = self.system_fmt.format(history[0]['content'])
+            else:
+                self.system = history[0]['content']
+            history = history[1:]
+        if history[0]['role'] != 'user':
+            raise ValueError('First message must be from user')
+        if history[-1]['role'] == 'user':
+            history.append({'role': 'assistant', 'content': ''})
+        return [[user['content'], assistant['content']] for user, assistant in zip(history[::2], history[1::2])]
+
+    def as_string(self, history: List[List[str]]) -> str:
+        """Returns the chat history as a string.
+
+        Args:
+            history: A list of lists of strings. Each inner list contains two strings: the user input and the assistant response.
+
+        Returns:
+            The chat history as a string, formatted in the provided syntax.
+        """
+        if history == []:
+            return ''
+        if isinstance(history[0], dict):
+            history = self.parse_structured_history(history)
+        text = self.system + ''.join([
+            self.turn_joiner.join([
+                self.user.format(item[0]),
+                self.assistant.format(item[1]),
+            ]) for item in history[:-1]
+        ])
+        text += self.user.format(history[-1][0])
+        text += self.response_prefix
+        return text
+
+    def format_response(self, response: str):
+        """Formats a response in the provided syntax.
+
+        It assumes that the prompt contained the prefix of the response as is common when training.
+
+        Args:
+            response: The response to format.
+
+        Returns:
+            The response, formatted in the provided syntax.
+        """
+        return response + self.response_suffix
+
+
+class ChatMLFormatter(ChatFormatter):
+    """A class for formatting the chat history in ChatML syntax.
+
+    Args:
+        system: The system prompt. This is an unformatted string that will be wrapped in ChatML syntax.
+        user: The user prompt. If None, a default ChatML value is used.
+        assistant: The assistant prompt. If None, a default ChatML value is used.
+    """
+
+    def __init__(self, system: str = None) -> None:
+        self.system_fmt = '<|im_start|>system\n{}<|im_end|>\n'
+        if system:
+            system = self.system_fmt.format(system)
+        super().__init__(system)
+
+
+class InstructFormatter(ChatFormatter):
+    """A class for formatting the chat history in MPTInstruct syntax.
+
+    Args:
+        system: The system prompt. If None, a default MPTInstruct-formatted prompt is used.
+        user: The user prompt. If None, a default MPTInstruct value is used.
+        assistant: The assistant prompt. If None, a default MPTInstruct value is used.
+    """
+
+    DEFAULT_SYSTEM_PROMPT = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
+
+    def __init__(self, system: str = None) -> None:
+        self.system_fmt = '{}\n'
+        if system:
+            system = self.system_fmt.format(system)
+        else:
+            system = self.system_fmt.format(self.DEFAULT_SYSTEM_PROMPT)
+
+        user ="\n### Instruction\n{}"
+        assistant = "\n### Response\n{}"
+        super().__init__(system, user, assistant, turn_joiner='\n')
+
+
+
+class Llama2ChatFormatter(ChatFormatter):
+    """A class for formatting the chat history in Llama2Chat syntax.
+
+    Args:
+        system: The system prompt. If None, a default Llama2Chat-formatted prompt is used.
+        user: The user prompt. If None, a default Llama2Chat value is used.
+        assistant: The assistant prompt. If None, a default Llama2Chat value is used.
+    """
+
+    DEFAULT_SYSTEM_PROMPT = """\
+You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information."""
+
+
+    def __init__(self, system: str = None) -> None:
+        self.system_fmt = "<s>[INST] <<SYS>>\n{}\n<</SYS>>\n\n"
+        if not system:
+            system = self.DEFAULT_SYSTEM_PROMPT
+        system = self.system_fmt.format(system)
+        user ="{}"
+        assistant = " [/INST] {} </s><s>[INST] "
+        super().__init__(system, user, assistant, turn_joiner='')

scripts/inference/hf_chat.py

@@ -12,27 +12,7 @@
                           PreTrainedModel, PreTrainedTokenizerBase,
                           StoppingCriteria, StoppingCriteriaList, TextStreamer)
 
-
-class ChatFormatter:
-    """A class for formatting the chat history.
-
-    Args:
-        system: The system prompt. If None, a default ChatML-formatted prompt is used.
-        user: The user prompt. If None, a default ChatML value is used.
-        assistant: The assistant prompt. If None, a default ChatML value is used.
-
-    Attributes:
-        system: The system prompt.
-        user: The user prompt.
-        assistant: The assistant prompt.
-        response_prefix: The response prefix (anything before {} in the assistant format string)
-    """
-
-    def __init__(self, system: str, user: str, assistant: str) -> None:
-        self.system = system if system else '<|im_start|>system\nA conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.<|im_end|>\n'
-        self.user = user if user else '<|im_start|>user\n{}<|im_end|>\n'
-        self.assistant = assistant if assistant else '<|im_start|>assistant\n{}<|im_end|>\n'
-        self.response_prefix = self.assistant.split('{}')[0]
+from llmfoundry.data.finetuning.chat import ChatFormatter, ChatMLFormatter, InstructFormatter, Llama2ChatFormatter
 
 
 class Conversation:
@@ -102,20 +82,9 @@ def __call__(self, input_ids: torch.LongTensor,
             "- Type 'quit' to end\n- Type 'system' to change the system prompt\n"
         )
 
-    def _history_as_formatted_str(self) -> str:
-        text = self.chat_format.system + ''.join([
-            '\n'.join([
-                self.chat_format.user.format(item[0]),
-                self.chat_format.assistant.format(item[1]),
-            ]) for item in self.history[:-1]
-        ])
-        text += self.chat_format.user.format(self.history[-1][0])
-        text += self.chat_format.response_prefix
-        return text
-
     def turn(self, user_inp: str) -> None:
         self.history.append([user_inp, ''])
-        conversation = self._history_as_formatted_str()
+        conversation = self.chat_format.as_string(self.history)
         input_ids = self.tokenizer(conversation, return_tensors='pt').input_ids
         input_ids = input_ids.to(self.model.device)
         # also stream to stdout
@@ -153,7 +122,7 @@ def __call__(self) -> None:
                 print(f'history: {self.history}')
                 continue
             elif user_inp == 'history_fmt':
-                print(f'history: {self._history_as_formatted_str()}')
+                print(f'history: {self.chat_format.as_string(self.history)}')
                 continue
             elif user_inp == 'system':
                 print('Enter a new system prompt:')
@@ -363,9 +332,23 @@ def main(args: Namespace) -> None:
         autocast_context = nullcontext()
         print('NOT using autocast...')
 
-    chat_format = ChatFormatter(system=args.system_prompt,
-                                user=args.user_msg_fmt,
-                                assistant=args.assistant_msg_fmt)
+    # Chat format
+    model_coarse_name = model.config.model_type
+    model_name = args.name_or_path
+    print(model_coarse_name, model_name)
+    if 'llama2' in model_coarse_name.lower():
+        chat_format = Llama2ChatFormatter(system=args.system_prompt)
+    elif model_name in ["mosaicml/mpt-7b-8k-instruct", "mosaicml/mpt-30b-instruct",  "mosaicml/mpt-7b-instruct"]:
+        chat_format = InstructFormatter(system=args.system_prompt)
+    elif model_name in ["mosaicml/mpt-7b-8k-chat", "mosaicml/mpt-30b-chat",  "mosaicml/mpt-7b-chat"]:
+        chat_format = ChatMLFormatter(system=args.system_prompt)
+    elif model_coarse_name == "mpt":
+        print("unrecognized model name, using MPT-Chat format")
+        chat_format = ChatMLFormatter(system=args.system_prompt)
+    else:
+        chat_format = ChatFormatter(system=args.system_prompt,
+                                    user=args.user_msg_fmt,
+                                    assistant=args.assistant_msg_fmt)
 
     conversation = Conversation(model=model,
                                 tokenizer=tokenizer,