fix bug

examples/lightrag_hf_demo.py

@@ -3,6 +3,7 @@ import sys
 
 from lightrag import LightRAG, QueryParam
 from lightrag.llm import hf_model_complete, hf_embedding
+from lightrag.utils import EmbeddingFunc
 from transformers import AutoModel,AutoTokenizer
 
 WORKING_DIR = "./dickens"
@@ -14,9 +15,13 @@ 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")
+    embedding_func=EmbeddingFunc(
+        tokenizer=AutoTokenizer.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
+    ),
 )
 
 

examples/lightrag_openai_demo.py

@@ -5,15 +5,15 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete
 from transformers import AutoModel,AutoTokenizer
 
-WORKING_DIR = "./dickens"
+WORKING_DIR = "/home/zrguo/code/myrag/agriculture"
 
 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
+    llm_model_func=gpt_4o_mini_complete
+    # llm_model_func=gpt_4o_complete
 )
 
 

lightrag/lightrag.py

@@ -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:openai_embedding)# 
+    embedding_func: EmbeddingFunc = field(default_factory=lambda:openai_embedding)
     embedding_batch_num: int = 32
     embedding_func_max_async: int = 16
 
@@ -103,13 +99,6 @@ 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}")
@@ -139,10 +128,9 @@ class LightRAG:
         self.chunk_entity_relation_graph = self.graph_storage_cls(
             namespace="chunk_entity_relation", global_config=asdict(self)
         )
+
         self.embedding_func = limit_async_func_call(self.embedding_func_max_async)(
-            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.embedding_func
         )
 
         self.entities_vdb = (