chiragjn commited on
Commit
0976781
1 Parent(s): 8a82d2e

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

Browse files
src/axolotl/utils/config/models/input/v0_4_1/__init__.py CHANGED
@@ -129,8 +129,10 @@ class RLType(str, Enum):
129
  class ChatTemplate(str, Enum):
130
  """Chat templates configuration subset"""
131
 
 
132
  chatml = "chatml" # pylint: disable=invalid-name
133
  inst = "inst" # pylint: disable=invalid-name
 
134
 
135
 
136
  class LoftQConfig(BaseModel):
@@ -518,7 +520,7 @@ class AxolotlInputConfig(
518
  ] = None
519
  gpu_memory_limit: Optional[Union[int, str]] = None
520
 
521
- chat_template: Optional[Union[Literal["chatml", "inst"], ChatTemplate]] = None
522
  default_system_message: Optional[str] = None
523
 
524
  # INTERNALS - document for now, generally not set externally
 
129
  class ChatTemplate(str, Enum):
130
  """Chat templates configuration subset"""
131
 
132
+ alpaca = "alpaca" # pylint: disable=invalid-name
133
  chatml = "chatml" # pylint: disable=invalid-name
134
  inst = "inst" # pylint: disable=invalid-name
135
+ gemma = "gemma" # pylint: disable=invalid-name
136
 
137
 
138
  class LoftQConfig(BaseModel):
 
520
  ] = None
521
  gpu_memory_limit: Optional[Union[int, str]] = None
522
 
523
+ chat_template: Optional[ChatTemplate] = None
524
  default_system_message: Optional[str] = None
525
 
526
  # INTERNALS - document for now, generally not set externally