Add vram_management_demo.py
This commit is contained in:
@@ -4,15 +4,15 @@ from lightrag import LightRAG, QueryParam
|
|||||||
from lightrag.llm import ollama_model_complete, ollama_embedding
|
from lightrag.llm import ollama_model_complete, ollama_embedding
|
||||||
from lightrag.utils import EmbeddingFunc
|
from lightrag.utils import EmbeddingFunc
|
||||||
|
|
||||||
# 工作目录和文本文件目录路径
|
# Working directory and the directory path for text files
|
||||||
WORKING_DIR = "./dickens"
|
WORKING_DIR = "./dickens"
|
||||||
TEXT_FILES_DIR = "/llm/mt"
|
TEXT_FILES_DIR = "/llm/mt"
|
||||||
|
|
||||||
# 如果工作目录不存在,则创建该目录
|
# Create the working directory if it doesn't exist
|
||||||
if not os.path.exists(WORKING_DIR):
|
if not os.path.exists(WORKING_DIR):
|
||||||
os.mkdir(WORKING_DIR)
|
os.mkdir(WORKING_DIR)
|
||||||
|
|
||||||
# 初始化 LightRAG
|
# Initialize LightRAG
|
||||||
rag = LightRAG(
|
rag = LightRAG(
|
||||||
working_dir=WORKING_DIR,
|
working_dir=WORKING_DIR,
|
||||||
llm_model_func=ollama_model_complete,
|
llm_model_func=ollama_model_complete,
|
||||||
@@ -24,7 +24,7 @@ rag = LightRAG(
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
# 读取 TEXT_FILES_DIR 目录下所有的 .txt 文件
|
# Read all .txt files from the TEXT_FILES_DIR directory
|
||||||
texts = []
|
texts = []
|
||||||
for filename in os.listdir(TEXT_FILES_DIR):
|
for filename in os.listdir(TEXT_FILES_DIR):
|
||||||
if filename.endswith('.txt'):
|
if filename.endswith('.txt'):
|
||||||
@@ -32,7 +32,7 @@ for filename in os.listdir(TEXT_FILES_DIR):
|
|||||||
with open(file_path, 'r', encoding='utf-8') as file:
|
with open(file_path, 'r', encoding='utf-8') as file:
|
||||||
texts.append(file.read())
|
texts.append(file.read())
|
||||||
|
|
||||||
# 批量插入文本到 LightRAG,带有重试机制
|
# Batch insert texts into LightRAG with a retry mechanism
|
||||||
def insert_texts_with_retry(rag, texts, retries=3, delay=5):
|
def insert_texts_with_retry(rag, texts, retries=3, delay=5):
|
||||||
for _ in range(retries):
|
for _ in range(retries):
|
||||||
try:
|
try:
|
||||||
@@ -45,7 +45,7 @@ def insert_texts_with_retry(rag, texts, retries=3, delay=5):
|
|||||||
|
|
||||||
insert_texts_with_retry(rag, texts)
|
insert_texts_with_retry(rag, texts)
|
||||||
|
|
||||||
# 执行不同类型的查询,并处理潜在的错误
|
# Perform different types of queries and handle potential errors
|
||||||
try:
|
try:
|
||||||
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="naive")))
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -66,12 +66,12 @@ try:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error performing hybrid search: {e}")
|
print(f"Error performing hybrid search: {e}")
|
||||||
|
|
||||||
# 清理 VRAM 资源的函数
|
# Function to clear VRAM resources
|
||||||
def clear_vram():
|
def clear_vram():
|
||||||
os.system("sudo nvidia-smi --gpu-reset")
|
os.system("sudo nvidia-smi --gpu-reset")
|
||||||
|
|
||||||
# 定期清理 VRAM 以防止溢出
|
# Regularly clear VRAM to prevent overflow
|
||||||
clear_vram_interval = 3600 # 每小时清理一次
|
clear_vram_interval = 3600 # Clear once every hour
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
@@ -79,4 +79,4 @@ while True:
|
|||||||
if current_time - start_time > clear_vram_interval:
|
if current_time - start_time > clear_vram_interval:
|
||||||
clear_vram()
|
clear_vram()
|
||||||
start_time = current_time
|
start_time = current_time
|
||||||
time.sleep(60) # 每分钟检查一次时间
|
time.sleep(60) # Check the time every minute
|
||||||
|
Reference in New Issue
Block a user