mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-02 03:26:17 +00:00
community: Add Baichuan Embeddings batch size (#22942)
- **Support batch size** Baichuan updates the document, indicating that up to 16 documents can be imported at a time - **Standardized model init arg names** - baichuan_api_key -> api_key - model_name -> model
This commit is contained in:
@@ -17,3 +17,13 @@ def test_baichuan_embedding_query() -> None:
|
||||
embedding = BaichuanTextEmbeddings() # type: ignore[call-arg]
|
||||
output = embedding.embed_query(document)
|
||||
assert len(output) == 1024 # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_baichuan_embeddings_multi_documents() -> None:
|
||||
"""Test Baichuan Text Embedding for documents with multi texts."""
|
||||
document = "午餐吃了螺蛳粉"
|
||||
doc_amount = 35
|
||||
embeddings = BaichuanTextEmbeddings() # type: ignore[call-arg]
|
||||
output = embeddings.embed_documents([document] * doc_amount)
|
||||
assert len(output) == doc_amount # type: ignore[arg-type]
|
||||
assert len(output[0]) == 1024 # type: ignore[index]
|
||||
|
18
libs/community/tests/unit_tests/embeddings/test_baichuan.py
Normal file
18
libs/community/tests/unit_tests/embeddings/test_baichuan.py
Normal file
@@ -0,0 +1,18 @@
|
||||
from typing import cast
|
||||
|
||||
from langchain_core.pydantic_v1 import SecretStr
|
||||
|
||||
from langchain_community.embeddings import BaichuanTextEmbeddings
|
||||
|
||||
|
||||
def test_sparkllm_initialization_by_alias() -> None:
|
||||
# Effective initialization
|
||||
embeddings = BaichuanTextEmbeddings( # type: ignore[call-arg]
|
||||
model="embedding_model", # type: ignore[arg-type]
|
||||
api_key="your-api-key", # type: ignore[arg-type]
|
||||
)
|
||||
assert embeddings.model_name == "embedding_model"
|
||||
assert (
|
||||
cast(SecretStr, embeddings.baichuan_api_key).get_secret_value()
|
||||
== "your-api-key"
|
||||
)
|
Reference in New Issue
Block a user