Added support for Baichuan LLM (#934)

### What problem does this PR solve?

- Added support for Baichuan LLM

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Co-authored-by: 海贼宅 <stu_xyx@163.com>
This commit is contained in:
yungongzi
2024-05-28 09:09:37 +08:00
committed by GitHub
parent ec6ae744a1
commit 9ffd7ae321
6 changed files with 169 additions and 3 deletions

View File

@@ -104,6 +104,15 @@ class OpenAIEmbed(Base):
return np.array(res.data[0].embedding), res.usage.total_tokens
class BaiChuanEmbed(OpenAIEmbed):
def __init__(self, key,
model_name='Baichuan-Text-Embedding',
base_url='https://api.baichuan-ai.com/v1'):
if not base_url:
base_url = "https://api.baichuan-ai.com/v1"
super().__init__(key, model_name, base_url)
class QWenEmbed(Base):
def __init__(self, key, model_name="text_embedding_v2", **kwargs):
dashscope.api_key = key