yujunjun/LightRAG
1
0
Fork 0
mirror of https://github.com/HKUDS/LightRAG.git synced 2025-10-30 17:29:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix bug

Browse Source
This commit is contained in:
LarFii 2024-10-15 21:11:12 +08:00
parent a92f7bfd61
commit 2190425d95
3 changed files with 14 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
examples/lightrag_hf_demo.py
View File

@ -3,6 +3,7 @@ import sys
from lightrag import LightRAG, QueryParam from lightrag import LightRAG, QueryParam
from lightrag.llm import hf_model_complete, hf_embedding from lightrag.llm import hf_model_complete, hf_embedding
from lightrag.utils import EmbeddingFunc
from transformers import AutoModel,AutoTokenizer from transformers import AutoModel,AutoTokenizer
WORKING_DIR = "./dickens" WORKING_DIR = "./dickens"
@ -14,9 +15,13 @@ rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=hf_model_complete, llm_model_func=hf_model_complete,
llm_model_name='meta-llama/Llama-3.1-8B-Instruct', llm_model_name='meta-llama/Llama-3.1-8B-Instruct',
embedding_func=hf_embedding, embedding_func=EmbeddingFunc(
tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"), tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
embedding_dim=384,
max_token_size=5000,
func=hf_embedding
),
) )

6
examples/lightrag_openai_demo.py
View File

@ -5,15 +5,15 @@ from lightrag import LightRAG, QueryParam
from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete
from transformers import AutoModel,AutoTokenizer from transformers import AutoModel,AutoTokenizer
WORKING_DIR = "./dickens" WORKING_DIR = "/home/zrguo/code/myrag/agriculture"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=gpt_4o_complete llm_model_func=gpt_4o_mini_complete
# llm_model_func=gpt_4o_mini_complete # llm_model_func=gpt_4o_complete
) )

18
lightrag/lightrag.py
View File

@ -76,12 +76,8 @@ class LightRAG:
} }
) )
# text embedding
tokenizer: Any = None
embed_model: Any = None
# embedding_func: EmbeddingFunc = field(default_factory=lambda:hf_embedding) # embedding_func: EmbeddingFunc = field(default_factory=lambda:hf_embedding)
embedding_func: EmbeddingFunc = field(default_factory=lambda:openai_embedding)# embedding_func: EmbeddingFunc = field(default_factory=lambda:openai_embedding)
embedding_batch_num: int = 32 embedding_batch_num: int = 32
embedding_func_max_async: int = 16 embedding_func_max_async: int = 16
@ -103,13 +99,6 @@ class LightRAG:
convert_response_to_json_func: callable = convert_response_to_json convert_response_to_json_func: callable = convert_response_to_json
def __post_init__(self): def __post_init__(self):
if callable(self.embedding_func) and self.embedding_func.__name__ == 'hf_embedding':
if self.tokenizer is None:
self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
if self.embed_model is None:
self.embed_model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
log_file = os.path.join(self.working_dir, "lightrag.log") log_file = os.path.join(self.working_dir, "lightrag.log")
set_logger(log_file) set_logger(log_file)
logger.info(f"Logger initialized for working directory: {self.working_dir}") logger.info(f"Logger initialized for working directory: {self.working_dir}")
@ -139,10 +128,9 @@ class LightRAG:
self.chunk_entity_relation_graph = self.graph_storage_cls( self.chunk_entity_relation_graph = self.graph_storage_cls(
namespace="chunk_entity_relation", global_config=asdict(self) namespace="chunk_entity_relation", global_config=asdict(self)
) )
self.embedding_func = limit_async_func_call(self.embedding_func_max_async)( self.embedding_func = limit_async_func_call(self.embedding_func_max_async)(
lambda texts: self.embedding_func(texts, self.tokenizer, self.embed_model) self.embedding_func
if callable(self.embedding_func) and self.embedding_func.__name__ == 'hf_embedding'
else self.embedding_func(texts)
) )
self.entities_vdb = ( self.entities_vdb = (