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_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings, BedrockEmbeddings
from langchain_google_genai import GoogleGenerativeAIEmbeddings from langchain_google_genai import GoogleGenerativeAIEmbeddings
from ..helpers import models_tokens 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): class AbstractGraph(ABC):
@ -200,6 +200,12 @@ class AbstractGraph(ABC):
elif "claude-3-" in llm_params["model"]: elif "claude-3-" in llm_params["model"]:
self.model_token = models_tokens["claude"]["claude3"] self.model_token = models_tokens["claude"]["claude3"]
return Anthropic(llm_params) 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: else:
raise ValueError( raise ValueError(
"Model provided by the configuration not supported") "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.9-llama3-8b-gguf": 8192,
"cognitivecomputations/dolphin-2.8-mistral-7b-v02": 32768, "cognitivecomputations/dolphin-2.8-mistral-7b-v02": 32768,
"cognitivecomputations/dolphin-2.5-mixtral-8x7b": 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 .bedrock import Bedrock
from .anthropic import Anthropic from .anthropic import Anthropic
from .claude import Claude 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)