mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-07-23 08:52:21 +00:00
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
import os
|
||
import time
|
||
from lightrag import LightRAG, QueryParam
|
||
from lightrag.llm import ollama_model_complete, ollama_embedding
|
||
from lightrag.utils import EmbeddingFunc
|
||
|
||
# 工作目录和文本文件目录路径
|
||
WORKING_DIR = "./dickens"
|
||
TEXT_FILES_DIR = "/llm/mt"
|
||
|
||
# 如果工作目录不存在，则创建该目录
|
||
if not os.path.exists(WORKING_DIR):
|
||
os.mkdir(WORKING_DIR)
|
||
|
||
# 初始化 LightRAG
|
||
rag = LightRAG(
|
||
working_dir=WORKING_DIR,
|
||
llm_model_func=ollama_model_complete,
|
||
llm_model_name="qwen2.5:3b-instruct-max-context",
|
||
embedding_func=EmbeddingFunc(
|
||
embedding_dim=768,
|
||
max_token_size=8192,
|
||
func=lambda texts: ollama_embedding(texts, embed_model="nomic-embed-text"),
|
||
),
|
||
)
|
||
|
||
# 读取 TEXT_FILES_DIR 目录下所有的 .txt 文件
|
||
texts = []
|
||
for filename in os.listdir(TEXT_FILES_DIR):
|
||
if filename.endswith('.txt'):
|
||
file_path = os.path.join(TEXT_FILES_DIR, filename)
|
||
with open(file_path, 'r', encoding='utf-8') as file:
|
||
texts.append(file.read())
|
||
|
||
# 批量插入文本到 LightRAG，带有重试机制
|
||
def insert_texts_with_retry(rag, texts, retries=3, delay=5):
|
||
for _ in range(retries):
|
||
try:
|
||
rag.insert(texts)
|
||
return
|
||
except Exception as e:
|
||
print(f"Error occurred during insertion: {e}. Retrying in {delay} seconds...")
|
||
time.sleep(delay)
|
||
raise RuntimeError("Failed to insert texts after multiple retries.")
|
||
|
||
insert_texts_with_retry(rag, texts)
|
||
|
||
# 执行不同类型的查询，并处理潜在的错误
|
||
try:
|
||
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")))
|
||
except Exception as e:
|
||
print(f"Error performing naive search: {e}")
|
||
|
||
try:
|
||
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="local")))
|
||
except Exception as e:
|
||
print(f"Error performing local search: {e}")
|
||
|
||
try:
|
||
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="global")))
|
||
except Exception as e:
|
||
print(f"Error performing global search: {e}")
|
||
|
||
try:
|
||
print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")))
|
||
except Exception as e:
|
||
print(f"Error performing hybrid search: {e}")
|
||
|
||
# 清理 VRAM 资源的函数
|
||
def clear_vram():
|
||
os.system("sudo nvidia-smi --gpu-reset")
|
||
|
||
# 定期清理 VRAM 以防止溢出
|
||
clear_vram_interval = 3600 # 每小时清理一次
|
||
start_time = time.time()
|
||
|
||
while True:
|
||
current_time = time.time()
|
||
if current_time - start_time > clear_vram_interval:
|
||
clear_vram()
|
||
start_time = current_time
|
||
time.sleep(60) # 每分钟检查一次时间 |