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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user