Update ChatTemplate enum to include alpaca and gemma (#1396)

This commit is contained in:
Chirag Jain
2024-03-13 20:36:02 +05:30
committed by GitHub
parent 8a82d2e0a4
commit 0976781e15

View File

@@ -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 AxolotlInputConfig(
] = 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