mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-11-21 04:23:52 +00:00
9923821d75
This parameter is no longer used. Its removal simplifies the API and clarifies that token length management is handled by upstream text chunking logic rather than the embedding wrapper.
70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
import sys
|
|
|
|
if sys.version_info < (3, 9):
|
|
pass
|
|
else:
|
|
pass
|
|
import pipmaster as pm # Pipmaster for dynamic library install
|
|
|
|
# install specific modules
|
|
if not pm.is_installed("lmdeploy"):
|
|
pm.install("lmdeploy")
|
|
|
|
from openai import (
|
|
APIConnectionError,
|
|
RateLimitError,
|
|
APITimeoutError,
|
|
)
|
|
from tenacity import (
|
|
retry,
|
|
stop_after_attempt,
|
|
wait_exponential,
|
|
retry_if_exception_type,
|
|
)
|
|
|
|
|
|
import numpy as np
|
|
import aiohttp
|
|
import base64
|
|
import struct
|
|
|
|
|
|
@retry(
|
|
stop=stop_after_attempt(3),
|
|
wait=wait_exponential(multiplier=1, min=4, max=60),
|
|
retry=retry_if_exception_type(
|
|
(RateLimitError, APIConnectionError, APITimeoutError)
|
|
),
|
|
)
|
|
async def siliconcloud_embedding(
|
|
texts: list[str],
|
|
model: str = "netease-youdao/bce-embedding-base_v1",
|
|
base_url: str = "https://api.siliconflow.cn/v1/embeddings",
|
|
max_token_size: int = 8192,
|
|
api_key: str = None,
|
|
) -> np.ndarray:
|
|
if api_key and not api_key.startswith("Bearer "):
|
|
api_key = "Bearer " + api_key
|
|
|
|
headers = {"Authorization": api_key, "Content-Type": "application/json"}
|
|
|
|
truncate_texts = [text[0:max_token_size] for text in texts]
|
|
|
|
payload = {"model": model, "input": truncate_texts, "encoding_format": "base64"}
|
|
|
|
base64_strings = []
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(base_url, headers=headers, json=payload) as response:
|
|
content = await response.json()
|
|
if "code" in content:
|
|
raise ValueError(content)
|
|
base64_strings = [item["embedding"] for item in content["data"]]
|
|
|
|
embeddings = []
|
|
for string in base64_strings:
|
|
decode_bytes = base64.b64decode(string)
|
|
n = len(decode_bytes) // 4
|
|
float_array = struct.unpack("<" + "f" * n, decode_bytes)
|
|
embeddings.append(float_array)
|
|
return np.array(embeddings)
|