put llama back

InternLM · Mar 7, 2024 · bff5c43 · bff5c43
1 parent c85fcf9
commit bff5c43
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/lmdeploy/model.py b/lmdeploy/model.py
@@ -57,6 +57,7 @@ def chat_template(self):
         return model
 
 
+@MODELS.register_module(name='llama')
 @MODELS.register_module(name='base')
 class BaseModel:
     """Base model."""

diff --git a/tests/test_lmdeploy/test_model.py b/tests/test_lmdeploy/test_model.py
@@ -43,7 +43,7 @@ def test_model_config(model_name, meta_instruction):
 
 
 def test_base_model():
-    model = MODELS.get('base')()
+    model = MODELS.get('llama')()
     assert model is not None
     assert model.capability == 'chat'
     assert model.get_prompt('test') == 'test'