From 82ed4611a59153bcb77d8eeb655f089de443ee5f Mon Sep 17 00:00:00 2001
From: Fangyin Cheng <staneyffer@gmail.com>
Date: Sat, 25 May 2024 17:26:04 +0800
Subject: [PATCH] feat(model): Support llama-3-sqlcoder-8b (#1559)

---
 dbgpt/configs/model_config.py     |  1 +
 dbgpt/model/adapter/hf_adapter.py | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/dbgpt/configs/model_config.py b/dbgpt/configs/model_config.py
index e56fbbd86..4632fe39e 100644
--- a/dbgpt/configs/model_config.py
+++ b/dbgpt/configs/model_config.py
@@ -197,6 +197,7 @@ LLM_MODEL_CONFIG = {
     "phi-3-small-8k-instruct": os.path.join(MODEL_PATH, "Phi-3-small-8k-instruct"),
     "phi-3-mini-128k-instruct": os.path.join(MODEL_PATH, "Phi-3-mini-128k-instruct"),
     "phi-3-mini-4k-instruct": os.path.join(MODEL_PATH, "Phi-3-mini-4k-instruct"),
+    "llama-3-sqlcoder-8b": os.path.join(MODEL_PATH, "llama-3-sqlcoder-8b"),
 }
 
 EMBEDDING_MODEL_CONFIG = {
diff --git a/dbgpt/model/adapter/hf_adapter.py b/dbgpt/model/adapter/hf_adapter.py
index 9e2aecee9..d4eefa99a 100644
--- a/dbgpt/model/adapter/hf_adapter.py
+++ b/dbgpt/model/adapter/hf_adapter.py
@@ -451,6 +451,19 @@ class PhiAdapter(NewHFChatModelAdapter):
         return str_prompt
 
 
+class SQLCoderAdapter(Llama3Adapter):
+    """
+    https://huggingface.co/defog/llama-3-sqlcoder-8b
+    """
+
+    def do_match(self, lower_model_name_or_path: Optional[str] = None):
+        return (
+            lower_model_name_or_path
+            and "llama-3" in lower_model_name_or_path
+            and "sqlcoder" in lower_model_name_or_path
+        )
+
+
 # The following code is used to register the model adapter
 # The last registered model adapter is matched first
 register_model_adapter(YiAdapter)
@@ -465,3 +478,4 @@ register_model_adapter(Llama3Adapter)
 register_model_adapter(DeepseekV2Adapter)
 register_model_adapter(SailorAdapter)
 register_model_adapter(PhiAdapter)
+register_model_adapter(SQLCoderAdapter)