Improvement[Embeddings] Add dimension support to ZhipuAIEmbeddings (#25274)

- In the in ` embedding-3 ` and later models of Zhipu AI, it is
supported to specify the dimensions parameter of Embedding. Ref:
https://bigmodel.cn/dev/api#text_embedding-3 .
- Add test case for `embedding-3` model by assigning dimensions.
This commit is contained in:
ZhangShenao
2024-08-12 04:20:37 +08:00
committed by GitHub
parent 9cd608efb3
commit 43deed2a95
2 changed files with 25 additions and 2 deletions

View File

@@ -18,3 +18,14 @@ def test_zhipuai_embedding_query() -> None:
embedding = ZhipuAIEmbeddings() # type: ignore[call-arg]
res = embedding.embed_query(document)
assert len(res) == 1024 # type: ignore[arg-type]
def test_zhipuai_embedding_dimensions() -> None:
"""Test ZhipuAI Text Embedding for query by assigning dimensions"""
document = "This is a test query."
embedding = ZhipuAIEmbeddings(
model="embedding-3",
dimensions=2048,
) # type: ignore[call-arg]
res = embedding.embed_query(document)
assert len(res) == 2048 # type: ignore[arg-type]