Add starling alpha

Cyrilvallez · Apr 15, 2024 · c7eb68a · c7eb68a
1 parent 270dfc7
commit c7eb68a
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 0 deletions.
diff --git a/textwiz/configs/causal/starling.py b/textwiz/configs/causal/starling.py
@@ -4,6 +4,7 @@
 
 # Mistral models
 MODELS_MAPPING = {
+    'starling-7B-alpha': 'berkeley-nest/Starling-LM-7B-alpha',
     'starling-7B-beta': 'Nexusflow/Starling-LM-7B-beta',
 }
 MODELS_DTYPES = {model: torch.bfloat16 for model in MODELS_MAPPING.keys()}
@@ -13,6 +14,7 @@
 
 
 MODELS_VERSIONS = {
+    'starling-7B-alpha': {'transformers': '>=4.35.0', 'tokenizers': '>=0.14.0'},
     'starling-7B-beta': {'transformers': '>=4.37.1', 'tokenizers': '>=0.14.0'},
 }
 MODELS_ADDITIONAL_TOKENIZER_KWARGS = {model: {'use_fast': False} for model in MODELS_MAPPING.keys()}
diff --git a/textwiz/templates/conversation_template.py b/textwiz/templates/conversation_template.py
@@ -632,6 +632,7 @@ def get_prompt(self) -> str:
     'zephyr-7B-beta': ZephyrConversation,
 
     # Starling
+    'starling-7B-alpha': StarlingConversation,
     'starling-7B-beta': StarlingConversation,
 
     # Command

diff --git a/textwiz/templates/prompt_template.py b/textwiz/templates/prompt_template.py
@@ -479,6 +479,7 @@ def format_chat(self, prompt: str, model_context: str = '', system_prompt: str =
     'zephyr-7B-beta': ZephyrPromptTemplate,
 
     # Starling
+    'starling-7B-alpha': StarlingPromptTemplate,
     'starling-7B-beta': StarlingPromptTemplate,
 
     # Command