From 756133512ee1898ef2e774a6bb0f2130dd2720e2 Mon Sep 17 00:00:00 2001
From: LarFii <834462287@qq.com>
Date: Tue, 15 Oct 2024 21:21:57 +0800
Subject: [PATCH] fix bug

---
 examples/lightrag_hf_demo.py     | 8 +++++---
 examples/lightrag_openai_demo.py | 2 +-
 lightrag/__init__.py             | 2 +-
 lightrag/llm.py                  | 5 -----
 4 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/examples/lightrag_hf_demo.py b/examples/lightrag_hf_demo.py
index 4cd503b3..baf62bdb 100644
--- a/examples/lightrag_hf_demo.py
+++ b/examples/lightrag_hf_demo.py
@@ -16,11 +16,13 @@ rag = LightRAG(
     llm_model_func=hf_model_complete,  
     llm_model_name='meta-llama/Llama-3.1-8B-Instruct',
     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
+        func=lambda texts: hf_embedding(
+            texts, 
+            tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
+            embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+        )
     ),
 )
 
diff --git a/examples/lightrag_openai_demo.py b/examples/lightrag_openai_demo.py
index 507b2eb2..fb1f055c 100644
--- a/examples/lightrag_openai_demo.py
+++ b/examples/lightrag_openai_demo.py
@@ -5,7 +5,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete
 from transformers import AutoModel,AutoTokenizer
 
-WORKING_DIR = "/home/zrguo/code/myrag/agriculture"
+WORKING_DIR = "./dickens"
 
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
diff --git a/lightrag/__init__.py b/lightrag/__init__.py
index b3d1d4ca..dc8faa6a 100644
--- a/lightrag/__init__.py
+++ b/lightrag/__init__.py
@@ -1,5 +1,5 @@
 from .lightrag import LightRAG, QueryParam
 
-__version__ = "0.0.4"
+__version__ = "0.0.5"
 __author__ = "Zirui Guo"
 __url__ = "https://github.com/HKUDS/LightRAG"
diff --git a/lightrag/llm.py b/lightrag/llm.py
index 87b156c5..bcb7e495 100644
--- a/lightrag/llm.py
+++ b/lightrag/llm.py
@@ -141,11 +141,6 @@ async def openai_embedding(texts: list[str]) -> np.ndarray:
     return np.array([dp.embedding for dp in response.data])
 
 
-
-@wrap_embedding_func_with_attrs(
-    embedding_dim=384,
-    max_token_size=5000,
-)
 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():