add additional roles

NousResearch · Sep 6, 2024 · d8f8025 · d8f8025
1 parent ecb5b37
commit d8f8025
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/nanotron/data/chat_tokenizer.py b/src/nanotron/data/chat_tokenizer.py
@@ -108,6 +108,8 @@ def _get_role(self, message: dict) -> Tuple[str, bool]:
  role = message["from"]
  if role == "gpt" or role == "assistant":
  return "assistant", False
+ elif role == "summary" or role == "author":
+ return role, False
  elif role == "human":
  return "user", True
  else: