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

fix bug

Browse Source
This commit is contained in:
LarFii
2024-10-15 21:11:12 +08:00
parent f2e8732c43
commit cc4af0d75f
3 changed files with 14 additions and 21 deletions
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.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
),
)

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 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
)