|
import os |
|
|
|
from lightrag import LightRAG, QueryParam |
|
from lightrag.llm import ollama_model_complete, ollama_embedding |
|
from lightrag.utils import EmbeddingFunc |
|
|
|
WORKING_DIR = "./dickens" |
|
|
|
if not os.path.exists(WORKING_DIR): |
|
os.mkdir(WORKING_DIR) |
|
|
|
rag = LightRAG( |
|
working_dir=WORKING_DIR, |
|
llm_model_func=ollama_model_complete, |
|
llm_model_name='your_model_name', |
|
embedding_func=EmbeddingFunc( |
|
embedding_dim=768, |
|
max_token_size=8192, |
|
func=lambda texts: ollama_embedding( |
|
texts, |
|
embed_model="nomic-embed-text" |
|
) |
|
), |
|
) |
|
|
|
|
|
with open("./book.txt") as f: |
|
rag.insert(f.read()) |
|
|
|
|
|
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))) |
|
|
|
|
|
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))) |
|
|
|
|
|
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))) |
|
|
|
|
|
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))) |
|
|