tpoisonooo
commited on
Commit
Β·
7b67ad7
1
Parent(s):
a7f6abf
feat(examples): support siliconcloud free API
Browse files- README.md +1 -0
- examples/lightrag_siliconcloud_demo.py +79 -0
- lightrag/llm.py +47 -1
- requirements.txt +2 -1
README.md
CHANGED
@@ -629,6 +629,7 @@ def extract_queries(file_path):
|
|
629 |
β βββ lightrag_ollama_demo.py
|
630 |
β βββ lightrag_openai_compatible_demo.py
|
631 |
β βββ lightrag_openai_demo.py
|
|
|
632 |
β βββ vram_management_demo.py
|
633 |
βββ lightrag
|
634 |
β βββ __init__.py
|
|
|
629 |
β βββ lightrag_ollama_demo.py
|
630 |
β βββ lightrag_openai_compatible_demo.py
|
631 |
β βββ lightrag_openai_demo.py
|
632 |
+
β βββ lightrag_siliconcloud_demo.py
|
633 |
β βββ vram_management_demo.py
|
634 |
βββ lightrag
|
635 |
β βββ __init__.py
|
examples/lightrag_siliconcloud_demo.py
ADDED
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import asyncio
|
3 |
+
from lightrag import LightRAG, QueryParam
|
4 |
+
from lightrag.llm import openai_complete_if_cache, siliconcloud_embedding
|
5 |
+
from lightrag.utils import EmbeddingFunc
|
6 |
+
import numpy as np
|
7 |
+
|
8 |
+
WORKING_DIR = "./dickens"
|
9 |
+
|
10 |
+
if not os.path.exists(WORKING_DIR):
|
11 |
+
os.mkdir(WORKING_DIR)
|
12 |
+
|
13 |
+
|
14 |
+
async def llm_model_func(
|
15 |
+
prompt, system_prompt=None, history_messages=[], **kwargs
|
16 |
+
) -> str:
|
17 |
+
return await openai_complete_if_cache(
|
18 |
+
"Qwen/Qwen2.5-7B-Instruct",
|
19 |
+
prompt,
|
20 |
+
system_prompt=system_prompt,
|
21 |
+
history_messages=history_messages,
|
22 |
+
api_key=os.getenv("UPSTAGE_API_KEY"),
|
23 |
+
base_url="https://api.siliconflow.cn/v1/",
|
24 |
+
**kwargs,
|
25 |
+
)
|
26 |
+
|
27 |
+
|
28 |
+
async def embedding_func(texts: list[str]) -> np.ndarray:
|
29 |
+
return await siliconcloud_embedding(
|
30 |
+
texts,
|
31 |
+
model="netease-youdao/bce-embedding-base_v1",
|
32 |
+
api_key=os.getenv("UPSTAGE_API_KEY"),
|
33 |
+
max_token_size=int(512 * 1.5)
|
34 |
+
)
|
35 |
+
|
36 |
+
|
37 |
+
# function test
|
38 |
+
async def test_funcs():
|
39 |
+
result = await llm_model_func("How are you?")
|
40 |
+
print("llm_model_func: ", result)
|
41 |
+
|
42 |
+
result = await embedding_func(["How are you?"])
|
43 |
+
print("embedding_func: ", result)
|
44 |
+
|
45 |
+
|
46 |
+
asyncio.run(test_funcs())
|
47 |
+
|
48 |
+
|
49 |
+
rag = LightRAG(
|
50 |
+
working_dir=WORKING_DIR,
|
51 |
+
llm_model_func=llm_model_func,
|
52 |
+
embedding_func=EmbeddingFunc(
|
53 |
+
embedding_dim=768, max_token_size=512, func=embedding_func
|
54 |
+
),
|
55 |
+
)
|
56 |
+
|
57 |
+
|
58 |
+
with open("./book.txt") as f:
|
59 |
+
rag.insert(f.read())
|
60 |
+
|
61 |
+
# Perform naive search
|
62 |
+
print(
|
63 |
+
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
|
64 |
+
)
|
65 |
+
|
66 |
+
# Perform local search
|
67 |
+
print(
|
68 |
+
rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
|
69 |
+
)
|
70 |
+
|
71 |
+
# Perform global search
|
72 |
+
print(
|
73 |
+
rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
|
74 |
+
)
|
75 |
+
|
76 |
+
# Perform hybrid search
|
77 |
+
print(
|
78 |
+
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
|
79 |
+
)
|
lightrag/llm.py
CHANGED
@@ -2,8 +2,11 @@ import os
|
|
2 |
import copy
|
3 |
import json
|
4 |
import aioboto3
|
|
|
5 |
import numpy as np
|
6 |
import ollama
|
|
|
|
|
7 |
from openai import AsyncOpenAI, APIConnectionError, RateLimitError, Timeout
|
8 |
from tenacity import (
|
9 |
retry,
|
@@ -312,7 +315,7 @@ async def ollama_model_complete(
|
|
312 |
@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
|
313 |
@retry(
|
314 |
stop=stop_after_attempt(3),
|
315 |
-
wait=wait_exponential(multiplier=1, min=4, max=
|
316 |
retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
|
317 |
)
|
318 |
async def openai_embedding(
|
@@ -332,6 +335,49 @@ async def openai_embedding(
|
|
332 |
)
|
333 |
return np.array([dp.embedding for dp in response.data])
|
334 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
335 |
|
336 |
# @wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
|
337 |
# @retry(
|
|
|
2 |
import copy
|
3 |
import json
|
4 |
import aioboto3
|
5 |
+
import aiohttp
|
6 |
import numpy as np
|
7 |
import ollama
|
8 |
+
import base64
|
9 |
+
import struct
|
10 |
from openai import AsyncOpenAI, APIConnectionError, RateLimitError, Timeout
|
11 |
from tenacity import (
|
12 |
retry,
|
|
|
315 |
@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
|
316 |
@retry(
|
317 |
stop=stop_after_attempt(3),
|
318 |
+
wait=wait_exponential(multiplier=1, min=4, max=60),
|
319 |
retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
|
320 |
)
|
321 |
async def openai_embedding(
|
|
|
335 |
)
|
336 |
return np.array([dp.embedding for dp in response.data])
|
337 |
|
338 |
+
@retry(
|
339 |
+
stop=stop_after_attempt(3),
|
340 |
+
wait=wait_exponential(multiplier=1, min=4, max=60),
|
341 |
+
retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
|
342 |
+
)
|
343 |
+
async def siliconcloud_embedding(
|
344 |
+
texts: list[str],
|
345 |
+
model: str = "netease-youdao/bce-embedding-base_v1",
|
346 |
+
base_url: str = "https://api.siliconflow.cn/v1/embeddings",
|
347 |
+
max_token_size: int = 512,
|
348 |
+
api_key: str = None,
|
349 |
+
) -> np.ndarray:
|
350 |
+
if api_key and not api_key.startswith('Bearer '):
|
351 |
+
api_key = 'Bearer ' + api_key
|
352 |
+
|
353 |
+
headers = {
|
354 |
+
"Authorization": api_key,
|
355 |
+
"Content-Type": "application/json"
|
356 |
+
}
|
357 |
+
|
358 |
+
truncate_texts = [text[0:max_token_size] for text in texts]
|
359 |
+
|
360 |
+
payload = {
|
361 |
+
"model": model,
|
362 |
+
"input": truncate_texts,
|
363 |
+
"encoding_format": "base64"
|
364 |
+
}
|
365 |
+
|
366 |
+
base64_strings = []
|
367 |
+
async with aiohttp.ClientSession() as session:
|
368 |
+
async with session.post(base_url, headers=headers, json=payload) as response:
|
369 |
+
content = await response.json()
|
370 |
+
if 'code' in content:
|
371 |
+
raise ValueError(content)
|
372 |
+
base64_strings = [item['embedding'] for item in content['data']]
|
373 |
+
|
374 |
+
embeddings = []
|
375 |
+
for string in base64_strings:
|
376 |
+
decode_bytes = base64.b64decode(string)
|
377 |
+
n = len(decode_bytes) // 4
|
378 |
+
float_array = struct.unpack('<' + 'f' * n, decode_bytes)
|
379 |
+
embeddings.append(float_array)
|
380 |
+
return np.array(embeddings)
|
381 |
|
382 |
# @wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
|
383 |
# @retry(
|
requirements.txt
CHANGED
@@ -11,4 +11,5 @@ tiktoken
|
|
11 |
torch
|
12 |
transformers
|
13 |
xxhash
|
14 |
-
pyvis
|
|
|
|
11 |
torch
|
12 |
transformers
|
13 |
xxhash
|
14 |
+
pyvis
|
15 |
+
aiohttp
|