From 0976781e150afabad4e21f42677fb02e1a969280 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 13 Mar 2024 20:36:02 +0530 Subject: [PATCH] Update ChatTemplate enum to include alpaca and gemma (#1396) --- src/axolotl/utils/config/models/input/v0_4_1/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py index d7e7b24de4..dfe9a9be96 100644 --- a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py +++ b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py @@ -129,8 +129,10 @@ class RLType(str, Enum): class ChatTemplate(str, Enum): """Chat templates configuration subset""" + alpaca = "alpaca" # pylint: disable=invalid-name chatml = "chatml" # pylint: disable=invalid-name inst = "inst" # pylint: disable=invalid-name + gemma = "gemma" # pylint: disable=invalid-name class LoftQConfig(BaseModel): @@ -518,7 +520,7 @@ class Config: ] = None gpu_memory_limit: Optional[Union[int, str]] = None - chat_template: Optional[Union[Literal["chatml", "inst"], ChatTemplate]] = None + chat_template: Optional[ChatTemplate] = None default_system_message: Optional[str] = None # INTERNALS - document for now, generally not set externally