gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'HKUDS:main' into main

Browse Source
This commit is contained in:
tmuife
2024-12-05 10:59:10 +08:00
committed by GitHub
parent ffc916a6d4 dff3fe18e9
commit 0744d30a60
8 changed files with 134 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
examples/insert_custom_kg.py
View File

@@ -56,18 +56,6 @@ custom_kg = {
"description": "An annual technology conference held in CityC",
"source_id": "Source3",
},
{
"entity_name": "CompanyD",
"entity_type": "Organization",
"description": "A financial services company specializing in insurance",
"source_id": "Source4",
},
{
"entity_name": "ServiceZ",
"entity_type": "Service",
"description": "An insurance product offered by CompanyD",
"source_id": "Source4",
},
],
"relationships": [
{
@@ -94,13 +82,23 @@ custom_kg = {
"weight": 0.8,
"source_id": "Source3",
},
],
"chunks": [
{
"src_id": "CompanyD",
"tgt_id": "ServiceZ",
"description": "CompanyD provides ServiceZ",
"keywords": "provide, offer",
"weight": 1.0,
"source_id": "Source4",
"content": "ProductX, developed by CompanyA, has revolutionized the market with its cutting-edge features.",
"source_id": "Source1",
},
{
"content": "PersonA is a prominent researcher at UniversityB, focusing on artificial intelligence and machine learning.",
"source_id": "Source2",
},
{
"content": "EventY, held in CityC, attracts technology enthusiasts and companies from around the globe.",
"source_id": "Source3",
},
{
"content": "None",
"source_id": "UNKNOWN",
},
],
}

59
examples/lightrag_nvidia_demo.py
View File

@@ -1,11 +1,14 @@
import os
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm import openai_complete_if_cache, nvidia_openai_embedding, nvidia_openai_complete
from lightrag.llm import (
openai_complete_if_cache,
nvidia_openai_embedding,
)
from lightrag.utils import EmbeddingFunc
import numpy as np
#for custom llm_model_func
# for custom llm_model_func
from lightrag.utils import locate_json_string_body_from_string
WORKING_DIR = "./dickens"
@@ -13,14 +16,15 @@ WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR)
#some method to use your API key (choose one)
# some method to use your API key (choose one)
# NVIDIA_OPENAI_API_KEY = os.getenv("NVIDIA_OPENAI_API_KEY")
NVIDIA_OPENAI_API_KEY = "nvapi-xxxx" #your api key
NVIDIA_OPENAI_API_KEY = "nvapi-xxxx" # your api key
# using pre-defined function for nvidia LLM API. OpenAI compatible
# llm_model_func = nvidia_openai_complete
#If you trying to make custom llm_model_func to use llm model on NVIDIA API like other example:
# If you trying to make custom llm_model_func to use llm model on NVIDIA API like other example:
async def llm_model_func(
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
@@ -37,36 +41,41 @@ async def llm_model_func(
return locate_json_string_body_from_string(result)
return result
#custom embedding
# custom embedding
nvidia_embed_model = "nvidia/nv-embedqa-e5-v5"
async def indexing_embedding_func(texts: list[str]) -> np.ndarray:
return await nvidia_openai_embedding(
texts,
model = nvidia_embed_model, #maximum 512 token
model=nvidia_embed_model, # maximum 512 token
# model="nvidia/llama-3.2-nv-embedqa-1b-v1",
api_key=NVIDIA_OPENAI_API_KEY,
base_url="https://integrate.api.nvidia.com/v1",
input_type = "passage",
trunc = "END", #handling on server side if input token is longer than maximum token
encode = "float"
input_type="passage",
trunc="END", # handling on server side if input token is longer than maximum token
encode="float",
)
async def query_embedding_func(texts: list[str]) -> np.ndarray:
return await nvidia_openai_embedding(
texts,
model = nvidia_embed_model, #maximum 512 token
model=nvidia_embed_model, # maximum 512 token
# model="nvidia/llama-3.2-nv-embedqa-1b-v1",
api_key=NVIDIA_OPENAI_API_KEY,
base_url="https://integrate.api.nvidia.com/v1",
input_type = "query",
trunc = "END", #handling on server side if input token is longer than maximum token
encode = "float"
input_type="query",
trunc="END", # handling on server side if input token is longer than maximum token
encode="float",
)
#dimension are same
# dimension are same
async def get_embedding_dim():
test_text = ["This is a test sentence."]
embedding = await indexing_embedding_func(test_text)
embedding = await indexing_embedding_func(test_text)
embedding_dim = embedding.shape[1]
return embedding_dim
@@ -88,29 +97,29 @@ async def main():
embedding_dimension = await get_embedding_dim()
print(f"Detected embedding dimension: {embedding_dimension}")
#lightRAG class during indexing
# lightRAG class during indexing
rag = LightRAG(
working_dir=WORKING_DIR,
llm_model_func=llm_model_func,
# llm_model_name="meta/llama3-70b-instruct", #un comment if
# llm_model_name="meta/llama3-70b-instruct", #un comment if
embedding_func=EmbeddingFunc(
embedding_dim=embedding_dimension,
max_token_size=512, #maximum token size, somehow it's still exceed maximum number of token
#so truncate (trunc) parameter on embedding_func will handle it and try to examine the tokenizer used in LightRAG
#so you can adjust to be able to fit the NVIDIA model (future work)
max_token_size=512, # maximum token size, somehow it's still exceed maximum number of token
# so truncate (trunc) parameter on embedding_func will handle it and try to examine the tokenizer used in LightRAG
# so you can adjust to be able to fit the NVIDIA model (future work)
func=indexing_embedding_func,
),
)
#reading file
# reading file
with open("./book.txt", "r", encoding="utf-8") as f:
await rag.ainsert(f.read())
#redefine rag to change embedding into query type
# redefine rag to change embedding into query type
rag = LightRAG(
working_dir=WORKING_DIR,
llm_model_func=llm_model_func,
# llm_model_name="meta/llama3-70b-instruct", #un comment if
# llm_model_name="meta/llama3-70b-instruct", #un comment if
embedding_func=EmbeddingFunc(
embedding_dim=embedding_dimension,
max_token_size=512,