use fastchat conversations template

README.md

@@ -180,7 +180,7 @@ Have dataset(s) in one of the following format (JSONL recommended):
   ```json
   {"instruction": "...", "input": "...", "output": "..."}
   ```
-- `sharegpt:chat`: conversations where `from` is `human`/`gpt`
+- `sharegpt`: conversations where `from` is `human`/`gpt`
   ```json
   {"conversations": [{"from": "...", "value": "..."}]}
   ```
@@ -269,11 +269,11 @@ Have dataset(s) in one of the following format (JSONL recommended):
   ```json
   {"prompt": "...", "generation": "..."}
   ```
-- `sharegpt_simple.load_role`: conversations where `role` is used instead of `from`
+- `sharegpt.load_role`: conversations where `role` is used instead of `from`
   ```json
   {"conversations": [{"role": "...", "value": "..."}]}
   ```
-- `sharegpt_simple.load_guanaco`: conversations where `from` is `prompter`/`assistant` instead of default sharegpt
+- `sharegpt.load_guanaco`: conversations where `from` is `prompter`/`assistant` instead of default sharegpt
   ```json
   {"conversations": [{"from": "...", "value": "..."}]}
   ```
@@ -443,6 +443,7 @@ datasets:
     data_files: # Optional[str] path to source data files
     shards: # Optional[int] number of shards to split data into
     name: # Optional[str] name of dataset configuration to load
+    conversation:  # Optional[str] fastchat conversation type, only used with type: sharegpt
 
   # custom user prompt
   - path: repo

requirements.txt

@@ -31,3 +31,4 @@ scipy
 scikit-learn==1.2.2
 pynvml
 art
+fschat==0.2.29

src/axolotl/monkeypatch/fastchat_conversation_turns.py

@@ -0,0 +1,174 @@
+"""
+monkeypatch to add a get_turns method
+"""
+
+import logging
+from typing import Generator, Tuple
+
+from fastchat.conversation import SeparatorStyle
+
+LOG = logging.getLogger("axolotl.monkeypatch.fastchat_conversation_turns")
+
+
+def get_prompt(self) -> str:
+    ret = ""
+    for role, msg in self.get_turns():
+        ret += role + msg
+    return ret
+
+
+def get_turns(  # pylint: disable=too-many-return-statements
+    self,
+) -> Generator[Tuple[str, str], None, None]:
+    """Get the prompt for generation."""
+    system_prompt = self.system_template.format(system_message=self.system_message)
+    if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+        yield "", system_prompt + self.sep
+        for role, message in self.messages:
+            if message:
+                yield role + ": ", message + self.sep
+            else:
+                yield role + ":", ""
+        return
+    if self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+        seps = [self.sep, self.sep2]
+        yield "", system_prompt + seps[0]
+        for i, (role, message) in enumerate(self.messages):
+            if message:
+                yield role + ": ", message + seps[i % 2]
+            else:
+                yield role + ":", ""
+        return
+    if self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+        yield "", system_prompt + self.sep
+        for role, message in self.messages:
+            if message:
+                yield role + ": ", message + self.sep
+            else:
+                yield role + ": ", ""  # must be end with a space
+        return
+    if self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+        yield "", "" if system_prompt == "" else system_prompt + self.sep
+        for role, message in self.messages:
+            if message:
+                yield role + "\n", message + self.sep
+            else:
+                yield role + "\n", ""
+        return
+    if self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+        yield "", system_prompt
+        for role, message in self.messages:
+            if message:
+                yield role, message + self.sep
+            else:
+                yield role, ""
+        return
+    if self.sep_style == SeparatorStyle.NO_COLON_TWO:
+        seps = [self.sep, self.sep2]
+        yield "", system_prompt
+        for i, (role, message) in enumerate(self.messages):
+            if message:
+                yield role, message + seps[i % 2]
+            else:
+                yield role, ""
+        return
+    if self.sep_style == SeparatorStyle.RWKV:
+        yield "", system_prompt
+        for i, (role, message) in enumerate(self.messages):
+            if message:
+                yield role + ": ", message.replace("\r\n", "\n").replace(
+                    "\n\n", "\n"
+                ) + "\n\n"
+            else:
+                yield role + ":", ""
+        return
+    if self.sep_style == SeparatorStyle.LLAMA2:
+        seps = [self.sep, self.sep2]
+        if self.system_message:
+            yield "", system_prompt
+        else:
+            yield "", "[INST] "
+        for i, (role, message) in enumerate(self.messages[1:]):
+            if message:
+                yield role + " ", message + seps[i % 2]
+            else:
+                yield role, ""
+        return
+    if self.sep_style == SeparatorStyle.CHATGLM:
+        # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+        # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+        round_add_n = 1 if self.name == "chatglm2" else 0
+        if system_prompt:
+            yield "", system_prompt + self.sep
+
+        for i, (role, message) in enumerate(self.messages):
+            if i % 2 == 0:
+                yield "", f"[Round {i//2 + round_add_n}]{self.sep}"
+
+            if message:
+                yield f"{role}：", f"{message}{self.sep}"
+            else:
+                yield f"{role}：", ""
+        return
+    if self.sep_style == SeparatorStyle.CHATML:
+        yield "", "" if system_prompt == "" else system_prompt + self.sep + "\n"
+        for role, message in self.messages:
+            if message:
+                yield role + "\n", message + self.sep + "\n"
+            else:
+                yield role + "\n", ""
+        return
+    if self.sep_style == SeparatorStyle.CHATINTERN:
+        # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+        seps = [self.sep, self.sep2]
+        yield "", system_prompt
+        for i, (role, message) in enumerate(self.messages):
+            prefix = "<s>" if i % 2 == 0 else ""
+            if message:
+                yield prefix + role + ":", message + seps[i % 2] + "\n"
+            else:
+                yield role + ":", ""
+        return
+    if self.sep_style == SeparatorStyle.DOLLY:
+        seps = [self.sep, self.sep2]
+        yield "", system_prompt
+        for i, (role, message) in enumerate(self.messages):
+            if message:
+                suffix = "\n\n" if i % 2 == 1 else ""
+                yield role + ":\n", message + seps[i % 2] + suffix
+            else:
+                yield role + ":\n", ""
+        return
+    if self.sep_style == SeparatorStyle.PHOENIX:
+        yield "", system_prompt
+        for role, message in self.messages:
+            if message:
+                yield role + ": ", "<s>" + message + "</s>"
+            else:
+                yield role + ": " + "<s>", ""
+        return
+    if self.sep_style == SeparatorStyle.ROBIN:
+        yield "", system_prompt + self.sep
+        for role, message in self.messages:
+            if message:
+                yield role + ":\n", message + self.sep
+            else:
+                yield role + ":\n", ""
+        return
+    if self.sep_style == SeparatorStyle.FALCON_CHAT:
+        if self.system_message:
+            yield "", system_prompt + self.sep
+        for role, message in self.messages:
+            if message:
+                yield role + ": ", message + self.sep
+            else:
+                yield role + ":", ""
+    else:
+        raise ValueError(f"Invalid style: {self.sep_style}")
+
+
+def add_get_turns_to_conversation():
+    import fastchat.conversation
+
+    fastchat.conversation.Conversation.get_turns = get_turns
+    fastchat.conversation.Conversation.get_prompt = get_prompt

src/axolotl/prompt_strategies/sharegpt_simple.py → src/axolotl/prompt_strategies/sharegpt.py

@@ -1,12 +1,35 @@
 """Module containing the SimpleShareGPTPromptTokenizingStrategy class"""
+from typing import Any, Dict, Optional
+
+from fastchat.conversation import Conversation, SeparatorStyle, register_conv_template
 
 from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import PromptStyle, ShareGPTPrompter
+from axolotl.prompters import ShareGPTPrompterV2
+
+register_conv_template(
+    Conversation(
+        name="chatml",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are a helpful assistant.",
+        roles=["<|im_start|>user", "<|im_start|>assistant"],
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>\n",
+    )
+)
 
 
-def load(tokenizer, cfg):
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
+    conversation = (
+        ds_cfg["conversation"] if ds_cfg and "conversation" in ds_cfg else None
+    )
+    field_human = ds_cfg["field_human"] if ds_cfg and "field_human" in ds_cfg else None
+    field_model = ds_cfg["field_model"] if ds_cfg and "field_model" in ds_cfg else None
     return SimpleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
+        ShareGPTPrompterV2(
+            conversation=conversation,
+            role_key_model=field_model,
+            role_key_human=field_human,
+        ),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,
@@ -15,7 +38,7 @@ def load(tokenizer, cfg):
 
 def load_role(tokenizer, cfg):
     return SimpleRoleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
+        ShareGPTPrompterV2(),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,
@@ -24,7 +47,7 @@ def load_role(tokenizer, cfg):
 
 def load_guanaco(tokenizer, cfg):
     return GuanacoShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
+        ShareGPTPrompterV2(),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,

src/axolotl/prompt_strategies/sharegpt_jokes.py

@@ -1,11 +1,11 @@
 """Module for Jokes prompts using sharegpt style """
 from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
-from axolotl.prompters import PromptStyle, ShareGPTPrompter
+from axolotl.prompters import ShareGPTPrompterV2
 
 
 def load(tokenizer, cfg):
     return SimpleJokesShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompter(PromptStyle.CHAT.value),
+        ShareGPTPrompterV2(),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,

src/axolotl/prompt_tokenizers.py

@@ -6,8 +6,12 @@
 import logging
 from typing import Dict, List, Tuple, Union
 
+from fastchat.conversation import Conversation
 from transformers import BatchEncoding, PreTrainedTokenizer
 
+from axolotl.monkeypatch.fastchat_conversation_turns import (
+    add_get_turns_to_conversation,
+)
 from axolotl.prompters import IGNORE_TOKEN_ID
 
 LOG = logging.getLogger("axolotl")
@@ -18,6 +22,8 @@
 LLAMA_DEFAULT_BOS_TOKEN = "<s>"  # nosec
 LLAMA_DEFAULT_UNK_TOKEN = "<unk>"  # nosec
 
+add_get_turns_to_conversation()
+
 
 class InvalidDataException(Exception):
     """
@@ -352,33 +358,36 @@ def tokenize_prompt(self, prompt):
         result, current_len = tokenize_prompt_default()
         user_token = self._get_user_token()
         assistant_token = self._get_assistant_token()
+        conversation: Conversation = (
+            self.prompter._conversation  # pylint: disable=protected-access
+        )
         try:
             for _, part in enumerate(
                 self.prompter.build_prompt(self.get_conversation_thread(prompt))
             ):
                 if isinstance(part, tuple):
-                    if part[0] == "USER:":
+                    if conversation.roles[0] in part[0]:
                         turn = part[0] + part[1] if not user_token else part[1]
                         # this is still the user query, we should
                         if not part[1].strip():
                             LOG.warning(f"user turn has empty text: {prompt}")
                         res = self._tokenize(
-                            turn.strip(),
+                            turn,
                             add_eos_token=False,
                             strip_bos_token=True,
                         )
                         if user_token:
                             res["input_ids"] = [user_token, *res["input_ids"]]
                         # everything from this is masked out from the labels
                         labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
-                    elif part[0] == "ASSISTANT:":
+                    elif conversation.roles[1] in part[0]:
                         # TODO label assistant token/tokens w/ IGNORE_TOKEN_ID
                         turn = part[0] + part[1] if not assistant_token else part[1]
                         # this should be the assistant response, should end with an eos token
                         if not part[1].strip():
                             LOG.warning(f"assistant turn has empty text: {prompt}")
                         res = self._tokenize(
-                            turn.strip(),
+                            turn,
                             add_eos_token=True,
                             strip_bos_token=True,
                         )
@@ -389,16 +398,17 @@ def tokenize_prompt(self, prompt):
                             ]
                         # not masked out from labels
                         labels = copy.deepcopy(res["input_ids"])
-                    elif part[0] == "SYSTEM:":
-                        part = part[1]  # Ignore the system role from preamble
+                    elif part[0] == "":
+                        turn = part[1]
                         # this is only ever the first part, should include the bos token and the user query
                         res = self._tokenize(
-                            part.strip(), add_eos_token=False, strip_bos_token=False
+                            turn, add_eos_token=False, strip_bos_token=False
                         )
                         # everything from this is masked out from the labels
                         labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
                     else:
                         LOG.warning(f"unhandled role: {part[0]}")
+                        continue
 
                 # pylint: disable=duplicate-code
                 result, current_len = parse_tokenized_to_result(