From ea126a7108135a6ecd7a3d52aa5f5620bd92dfaa Mon Sep 17 00:00:00 2001 From: LarFii <834462287@qq.com> Date: Tue, 15 Oct 2024 19:40:08 +0800 Subject: [PATCH] Add huggingface model support --- README.md | 8 +++---- examples/insert.py | 18 ---------------- examples/lightrag_hf_demo.py | 36 ++++++++++++++++++++++++++++++++ examples/lightrag_openai_demo.py | 33 +++++++++++++++++++++++++++++ examples/query.py | 16 -------------- lightrag/__init__.py | 2 +- lightrag/base.py | 2 +- lightrag/lightrag.py | 29 ++++++++++++++++++------- lightrag/llm.py | 8 ++----- lightrag/operate.py | 2 +- reproduce/Step_3.py | 2 +- 11 files changed, 100 insertions(+), 56 deletions(-) delete mode 100644 examples/insert.py create mode 100644 examples/lightrag_hf_demo.py create mode 100644 examples/lightrag_openai_demo.py delete mode 100644 examples/query.py diff --git a/README.md b/README.md index 1541e153..562585e4 100644 --- a/README.md +++ b/README.md @@ -59,8 +59,8 @@ print(rag.query("What are the top themes in this story?", param=QueryParam(mode= # Perform global search print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))) -# Perform hybird search -print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybird"))) +# Perform hybrid search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))) ``` Batch Insert ```python @@ -287,8 +287,8 @@ def extract_queries(file_path): ├── examples │ ├── batch_eval.py │ ├── generate_query.py -│ ├── insert.py -│ └── query.py +│ ├── lightrag_openai_demo.py +│ └── lightrag_hf_demo.py ├── lightrag │ ├── __init__.py │ ├── base.py diff --git a/examples/insert.py b/examples/insert.py deleted file mode 100644 index 25c3cdda..00000000 --- a/examples/insert.py +++ /dev/null @@ -1,18 +0,0 @@ -import os -import sys - -from lightrag import LightRAG - -# os.environ["OPENAI_API_KEY"] = "" - -WORKING_DIR = "" - -if not os.path.exists(WORKING_DIR): - os.mkdir(WORKING_DIR) - -rag = LightRAG(working_dir=WORKING_DIR) - -with open('./text.txt', 'r') as f: - text = f.read() - -rag.insert(text) \ No newline at end of file diff --git a/examples/lightrag_hf_demo.py b/examples/lightrag_hf_demo.py new file mode 100644 index 00000000..f0e5fa99 --- /dev/null +++ b/examples/lightrag_hf_demo.py @@ -0,0 +1,36 @@ +import os +import sys + +from lightrag import LightRAG, QueryParam +from lightrag.llm import hf_model_complete, hf_embedding +from transformers import AutoModel,AutoTokenizer + +WORKING_DIR = "./dickens" + +if not os.path.exists(WORKING_DIR): + os.mkdir(WORKING_DIR) + +rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=hf_model_complete, + llm_model_name='meta-llama/Llama-3.1-8B-Instruct', + embedding_func=hf_embedding, + tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"), + embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") +) + + +with open("./book.txt") as f: + rag.insert(f.read()) + +# Perform naive search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))) + +# Perform local search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))) + +# Perform global search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))) + +# Perform hybrid search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))) diff --git a/examples/lightrag_openai_demo.py b/examples/lightrag_openai_demo.py new file mode 100644 index 00000000..677506c2 --- /dev/null +++ b/examples/lightrag_openai_demo.py @@ -0,0 +1,33 @@ +import os +import sys + +from lightrag import LightRAG, QueryParam +from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete +from transformers import AutoModel,AutoTokenizer + +WORKING_DIR = "./dickens" + +if not os.path.exists(WORKING_DIR): + os.mkdir(WORKING_DIR) + +rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=gpt_4o_complete + # llm_model_func=gpt_4o_mini_complete +) + + +with open("./book.txt") as f: + rag.insert(f.read()) + +# Perform naive search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))) + +# Perform local search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))) + +# Perform global search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))) + +# Perform hybrid search +print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))) diff --git a/examples/query.py b/examples/query.py deleted file mode 100644 index 00c902eb..00000000 --- a/examples/query.py +++ /dev/null @@ -1,16 +0,0 @@ -import os -import sys - -from lightrag import LightRAG, QueryParam - -# os.environ["OPENAI_API_KEY"] = "" - -WORKING_DIR = "" - -rag = LightRAG(working_dir=WORKING_DIR) - -mode = 'global' -query_param = QueryParam(mode=mode) - -result = rag.query("", param=query_param) -print(result) \ No newline at end of file diff --git a/lightrag/__init__.py b/lightrag/__init__.py index 0b279096..b3d1d4ca 100644 --- a/lightrag/__init__.py +++ b/lightrag/__init__.py @@ -1,5 +1,5 @@ from .lightrag import LightRAG, QueryParam -__version__ = "0.0.3" +__version__ = "0.0.4" __author__ = "Zirui Guo" __url__ = "https://github.com/HKUDS/LightRAG" diff --git a/lightrag/base.py b/lightrag/base.py index 9c0422fe..d677c406 100644 --- a/lightrag/base.py +++ b/lightrag/base.py @@ -14,7 +14,7 @@ T = TypeVar("T") @dataclass class QueryParam: - mode: Literal["local", "global", "hybird", "naive"] = "global" + mode: Literal["local", "global", "hybrid", "naive"] = "global" only_need_context: bool = False response_type: str = "Multiple Paragraphs" top_k: int = 60 diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 9c34a607..329bfd12 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -3,7 +3,8 @@ import os from dataclasses import asdict, dataclass, field from datetime import datetime from functools import partial -from typing import Type, cast +from typing import Type, cast, Any +from transformers import AutoModel,AutoTokenizer, AutoModelForCausalLM from .llm import gpt_4o_complete, gpt_4o_mini_complete, openai_embedding,hf_model_complete,hf_embedding from .operate import ( @@ -11,7 +12,7 @@ from .operate import ( extract_entities, local_query, global_query, - hybird_query, + hybrid_query, naive_query, ) @@ -38,15 +39,14 @@ from .base import ( def always_get_an_event_loop() -> asyncio.AbstractEventLoop: try: - # If there is already an event loop, use it. - loop = asyncio.get_event_loop() + loop = asyncio.get_running_loop() except RuntimeError: - # If in a sub-thread, create a new event loop. logger.info("Creating a new event loop in a sub-thread.") loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) return loop + @dataclass class LightRAG: working_dir: str = field( @@ -77,6 +77,9 @@ 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:openai_embedding)# embedding_batch_num: int = 32 @@ -100,6 +103,13 @@ class LightRAG: convert_response_to_json_func: callable = convert_response_to_json 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") set_logger(log_file) logger.info(f"Logger initialized for working directory: {self.working_dir}") @@ -130,8 +140,11 @@ class LightRAG: namespace="chunk_entity_relation", global_config=asdict(self) ) self.embedding_func = limit_async_func_call(self.embedding_func_max_async)( - self.embedding_func + lambda texts: self.embedding_func(texts, self.tokenizer, self.embed_model) + if callable(self.embedding_func) and self.embedding_func.__name__ == 'hf_embedding' + else self.embedding_func(texts) ) + self.entities_vdb = ( self.vector_db_storage_cls( namespace="entities", @@ -267,8 +280,8 @@ class LightRAG: param, asdict(self), ) - elif param.mode == "hybird": - response = await hybird_query( + elif param.mode == "hybrid": + response = await hybrid_query( query, self.chunk_entity_relation_graph, self.entities_vdb, diff --git a/lightrag/llm.py b/lightrag/llm.py index 5fb27b04..bc2ac1f3 100644 --- a/lightrag/llm.py +++ b/lightrag/llm.py @@ -142,18 +142,14 @@ async def openai_embedding(texts: list[str]) -> np.ndarray: -global EMBED_MODEL -global tokenizer -EMBED_MODEL = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") -tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") @wrap_embedding_func_with_attrs( embedding_dim=384, max_token_size=5000, ) -async def hf_embedding(texts: list[str]) -> np.ndarray: +async def hf_embedding(texts: list[str], tokenizer, embed_model) -> np.ndarray: input_ids = tokenizer(texts, return_tensors='pt', padding=True, truncation=True).input_ids with torch.no_grad(): - outputs = EMBED_MODEL(input_ids) + outputs = embed_model(input_ids) embeddings = outputs.last_hidden_state.mean(dim=1) return embeddings.detach().numpy() diff --git a/lightrag/operate.py b/lightrag/operate.py index a8213a37..3d388cb6 100644 --- a/lightrag/operate.py +++ b/lightrag/operate.py @@ -827,7 +827,7 @@ async def _find_related_text_unit_from_relationships( return all_text_units -async def hybird_query( +async def hybrid_query( query, knowledge_graph_inst: BaseGraphStorage, entities_vdb: BaseVectorStorage, diff --git a/reproduce/Step_3.py b/reproduce/Step_3.py index f7f7ee30..e97e2af6 100644 --- a/reproduce/Step_3.py +++ b/reproduce/Step_3.py @@ -52,7 +52,7 @@ def run_queries_and_save_to_json(queries, rag_instance, query_param, output_file if __name__ == "__main__": cls = "agriculture" - mode = "hybird" + mode = "hybrid" WORKING_DIR = "../{cls}" rag = LightRAG(working_dir=WORKING_DIR)