feat(model): Support Yi 1.5 16k models (eosphoros-ai#1543)

Aries-ckt · May 20, 2024 · 140d1ce · 140d1ce
1 parent d81bbec
commit 140d1ce
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/dbgpt/configs/model_config.py b/dbgpt/configs/model_config.py
@@ -178,7 +178,9 @@ def get_device() -> str:
     # https://huggingface.co/01-ai/Yi-1.5-6B-Chat
     "yi-1.5-6b-chat": os.path.join(MODEL_PATH, "Yi-1.5-6B-Chat"),
     "yi-1.5-9b-chat": os.path.join(MODEL_PATH, "Yi-1.5-9B-Chat"),
+    "yi-1.5-9b-chat-16k": os.path.join(MODEL_PATH, "Yi-1.5-9B-Chat-16K"),
     "yi-1.5-34b-chat": os.path.join(MODEL_PATH, "Yi-1.5-34B-Chat"),
+    "yi-1.5-34b-chat-16k": os.path.join(MODEL_PATH, "Yi-1.5-34B-Chat-16K"),
     # https://huggingface.co/google/gemma-7b-it
     "gemma-7b-it": os.path.join(MODEL_PATH, "gemma-7b-it"),
     # https://huggingface.co/google/gemma-2b-it