feat: add support for deepseek-chat

closes #222
This commit is contained in:
Federico Aguzzi 2024-05-12 01:00:11 +02:00
parent 3b9ec9b1f5
commit 156b67b91e
4 changed files with 33 additions and 2 deletions

View File

@ -7,7 +7,7 @@ from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings, BedrockEmbeddings
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from ..helpers import models_tokens
from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, Claude
from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, Claude, DeepSeek
class AbstractGraph(ABC):
@ -200,6 +200,12 @@ class AbstractGraph(ABC):
elif "claude-3-" in llm_params["model"]:
self.model_token = models_tokens["claude"]["claude3"]
return Anthropic(llm_params)
elif "deepseek" in llm_params["model"]:
try:
self.model_token = models_tokens["deepseek"][llm_params["model"]]
except KeyError as exc:
raise KeyError("Model not supported") from exc
return DeepSeek(llm_params)
else:
raise ValueError(
"Model provided by the configuration not supported")

View File

@ -102,6 +102,12 @@ models_tokens = {
"cognitivecomputations/dolphin-2.9-llama3-8b-gguf": 8192,
"cognitivecomputations/dolphin-2.8-mistral-7b-v02": 32768,
"cognitivecomputations/dolphin-2.5-mixtral-8x7b": 32768,
"TheBloke/dolphin-2.7-mixtral-8x7b-GGUF": 32768
"TheBloke/dolphin-2.7-mixtral-8x7b-GGUF": 32768,
"deepseek-ai/DeepSeek-V2": 131072,
"deepseek-ai/DeepSeek-V2-Chat": 131072
},
"deepseek": {
"deepseek-chat": 32768,
"deepseek-coder": 16384
}
}

View File

@ -13,3 +13,4 @@ from .groq import Groq
from .bedrock import Bedrock
from .anthropic import Anthropic
from .claude import Claude
from .deepseek import DeepSeek

View File

@ -0,0 +1,18 @@
"""
DeepSeek Module
"""
from langchain_openai import ChatOpenAI
class DeepSeek(ChatOpenAI):
"""
A wrapper for the ChatOpenAI class (DeepSeek uses an OpenAI-like API) that
provides default configuration and could be extended with additional methods
if needed.
Args:
llm_config (dict): Configuration parameters for the language model.
"""
def __init__(self, llm_config: dict):
super().__init__(**llm_config)