From 99993e5026beb68d918acdcb83015404675f2ef8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BB=84=E8=85=BE?= <101850389+hangters@users.noreply.github.com> Date: Thu, 29 Aug 2024 16:14:49 +0800 Subject: [PATCH] add support for Voyage AI (#2159) ### What problem does this PR solve? #1853 #2138 add support for Voyage AI ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Co-authored-by: Zhedong Cen Co-authored-by: Kevin Hu --- conf/llm_factories.json | 64 ++++++++++++++++++- rag/llm/__init__.py | 8 ++- rag/llm/embedding_model.py | 21 ++++++ rag/llm/rerank_model.py | 16 +++++ requirements.txt | 1 + requirements_arm.txt | 1 + .../svg/llm/{lepton.svg => lepton-ai.svg} | 0 web/src/assets/svg/llm/voyage.svg | 32 ++++++++++ .../user-setting/setting-model/constant.ts | 3 +- 9 files changed, 141 insertions(+), 5 deletions(-) rename web/src/assets/svg/llm/{lepton.svg => lepton-ai.svg} (100%) create mode 100644 web/src/assets/svg/llm/voyage.svg diff --git a/conf/llm_factories.json b/conf/llm_factories.json index bd0ac30ab..698443efd 100644 --- a/conf/llm_factories.json +++ b/conf/llm_factories.json @@ -2400,7 +2400,7 @@ ] }, { - "name": "Lepton", + "name": "LeptonAI", "logo": "", "tags": "LLM", "status": "1", @@ -3290,6 +3290,68 @@ "model_type": "chat" } ] + }, + { + "name": "Voyage AI", + "logo": "", + "tags": "TEXT EMBEDDING, TEXT RE-RANK", + "status": "1", + "llm": [ + { + "llm_name": "voyage-large-2-instruct", + "tags": "TEXT EMBEDDING,16000", + "max_tokens": 16000, + "model_type": "embedding" + }, + { + "llm_name": "voyage-finance-2", + "tags": "TEXT EMBEDDING,32000", + "max_tokens": 32000, + "model_type": "embedding" + }, + { + "llm_name": "voyage-multilingual-2", + "tags": "TEXT EMBEDDING,32000", + "max_tokens": 32000, + "model_type": "embedding" + }, + { + "llm_name": "voyage-law-2", + "tags": "TEXT EMBEDDING,16000", + "max_tokens": 16000, + "model_type": "embedding" + }, + { + "llm_name": "voyage-code-2", + "tags": "TEXT EMBEDDING,16000", + "max_tokens": 16000, + "model_type": "embedding" + }, + { + "llm_name": "voyage-large-2", + "tags": "TEXT EMBEDDING,16000", + "max_tokens": 16000, + "model_type": "embedding" + }, + { + "llm_name": "voyage-2", + "tags": "TEXT EMBEDDING,4000", + "max_tokens": 4000, + "model_type": "embedding" + }, + { + "llm_name": "rerank-1", + "tags": "RE-RANK, 8000", + "max_tokens": 8000, + "model_type": "rerank" + }, + { + "llm_name": "rerank-lite-1", + "tags": "RE-RANK, 4000", + "max_tokens": 4000, + "model_type": "rerank" + } + ] } ] } diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py index ef37d6446..b503b1762 100644 --- a/rag/llm/__init__.py +++ b/rag/llm/__init__.py @@ -45,7 +45,8 @@ EmbeddingModel = { "Upstage": UpstageEmbed, "SILICONFLOW": SILICONFLOWEmbed, "Replicate": ReplicateEmbed, - "BaiduYiyan": BaiduYiyanEmbed + "BaiduYiyan": BaiduYiyanEmbed, + "Voyage AI": VoyageEmbed } @@ -105,7 +106,7 @@ ChatModel = { "Tencent Hunyuan": HunyuanChat, "XunFei Spark": SparkChat, "BaiduYiyan": BaiduYiyanChat, - "Anthropic": AnthropicChat + "Anthropic": VoyageChat, } @@ -120,7 +121,8 @@ RerankModel = { "cohere": CoHereRerank, "TogetherAI": TogetherAIRerank, "SILICONFLOW": SILICONFLOWRerank, - "BaiduYiyan": BaiduYiyanRerank + "BaiduYiyan": BaiduYiyanRerank, + "Voyage AI": VoyageRerank } diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index bc8241372..257a32891 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -623,3 +623,24 @@ class BaiduYiyanEmbed(Base): np.array([r["embedding"] for r in res["data"]]), res["usage"]["total_tokens"], ) + + +class VoyageEmbed(Base): + def __init__(self, key, model_name, base_url=None): + import voyageai + + self.client = voyageai.Client(api_key=key) + self.model_name = model_name + + def encode(self, texts: list, batch_size=32): + res = self.client.embed( + texts=texts, model=self.model_name, input_type="document" + ) + return np.array(res.embeddings), res.total_tokens + + def encode_queries(self, text): + res = self.client.embed + res = self.client.embed( + texts=text, model=self.model_name, input_type="query" + ) + return np.array(res.embeddings), res.total_tokens diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index cc7f5525c..3a2a91092 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -311,3 +311,19 @@ class BaiduYiyanRerank(Base): rank = np.array([d["relevance_score"] for d in res["results"]]) indexs = [d["index"] for d in res["results"]] return rank[indexs], res["usage"]["total_tokens"] + + +class VoyageRerank(Base): + def __init__(self, key, model_name, base_url=None): + import voyageai + + self.client = voyageai.Client(api_key=key) + self.model_name = model_name + + def similarity(self, query: str, texts: list): + res = self.client.rerank( + query=query, documents=texts, model=self.model_name, top_k=len(texts) + ) + rank = np.array([r.relevance_score for r in res.results]) + indexs = [r.index for r in res.results] + return rank[indexs], res.total_tokens diff --git a/requirements.txt b/requirements.txt index 0bcd697af..0f2136606 100644 --- a/requirements.txt +++ b/requirements.txt @@ -86,6 +86,7 @@ torch==2.3.0 transformers==4.38.1 umap==0.1.1 volcengine==1.0.146 +voyageai==0.2.3 webdriver_manager==4.0.1 Werkzeug==3.0.3 wikipedia==1.4.0 diff --git a/requirements_arm.txt b/requirements_arm.txt index 1207d6d8d..01b9ce0e3 100644 --- a/requirements_arm.txt +++ b/requirements_arm.txt @@ -141,6 +141,7 @@ loguru==0.7.2 umap-learn fasttext==0.9.2 volcengine==1.0.141 +voyageai==0.2.3 opencv-python-headless==4.9.0.80 readability-lxml==0.8.1 html_text==0.6.2 diff --git a/web/src/assets/svg/llm/lepton.svg b/web/src/assets/svg/llm/lepton-ai.svg similarity index 100% rename from web/src/assets/svg/llm/lepton.svg rename to web/src/assets/svg/llm/lepton-ai.svg diff --git a/web/src/assets/svg/llm/voyage.svg b/web/src/assets/svg/llm/voyage.svg new file mode 100644 index 000000000..88ffbdff9 --- /dev/null +++ b/web/src/assets/svg/llm/voyage.svg @@ -0,0 +1,32 @@ + + + + + + + + + + + + \ No newline at end of file diff --git a/web/src/pages/user-setting/setting-model/constant.ts b/web/src/pages/user-setting/setting-model/constant.ts index 33cf76d86..a644c21c8 100644 --- a/web/src/pages/user-setting/setting-model/constant.ts +++ b/web/src/pages/user-setting/setting-model/constant.ts @@ -24,7 +24,7 @@ export const IconMap = { 'LM-Studio': 'lm-studio', 'OpenAI-API-Compatible': 'openai-api', cohere: 'cohere', - Lepton: 'lepton', + LeptonAI: 'lepton-ai', TogetherAI: 'together-ai', PerfXCloud: 'perfx-cloud', Upstage: 'upstage', @@ -38,6 +38,7 @@ export const IconMap = { 'Fish Audio': 'fish-audio', 'Tencent Cloud': 'tencent-cloud', Anthropic: 'anthropic', + 'Voyage AI': 'voyage', }; export const BedrockRegionList = [