import os from lightrag import LightRAG, QueryParam from lightrag.llm import gpt_4o_mini_complete, gpt_4o_complete
######### # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert() # import nest_asyncio # nest_asyncio.apply() #########
WORKING_DIR = "./books"
if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR)
rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=gpt_4o_mini_complete# Use gpt_4o_mini_complete LLM model # llm_model_func=gpt_4o_complete# Optionally, use a stronger model )
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")))
当然,你也可以选择本地model, 比如使用Hugging Face的模型:
from lightrag.llm import hf_model_complete, hf_embedding from transformers import AutoModel, AutoTokenizer
# Initialize LightRAG with Hugging Face model rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=hf_model_complete,# Use Hugging Face model for text generation llm_model_name='meta-llama/Llama-3.1-8B-Instruct',# Model name from Hugging Face # Use Hugging Face embedding function embedding_func=EmbeddingFunc( embedding_dim=384, max_token_size=5000, 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") ) ), )
或者使用Ollama模型:
from lightrag.llm import ollama_model_complete, ollama_embedding
# Initialize LightRAG with Ollama model rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=ollama_model_complete,# Use Ollama model for text generation llm_model_name='your_model_name', # Your model name # Use Ollama embedding function embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, func=lambda texts: ollama_embedding( texts, embed_model="nomic-embed-text" ) ), )