You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

68 lines
2.1 KiB

import asyncio
import os
import shutil
from raganything import RAGAnything, RAGAnythingConfig
from Util.RagUtil import create_llm_model_func, create_vision_model_func, create_embedding_func
import logging
# 在程序开始时添加以下配置
logging.basicConfig(
level=logging.INFO, # 设置日志级别为INFO
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# 或者如果你想更详细地控制日志输出
logger = logging.getLogger('lightrag')
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)
async def main():
# 要处理的文件路径
file_path = "static/Txt/吉林动画学院一览表.pdf"
# 索引生成目录
WORKING_DIR = "./Topic/DongHua"
# 删除output目录下的所有文件
output_dir = "./output"
shutil.rmtree(output_dir, ignore_errors=True)
os.makedirs(output_dir, exist_ok=True)
# 删除WORKING_DIR下的所有文件
shutil.rmtree(WORKING_DIR, ignore_errors=True)
os.makedirs(WORKING_DIR, exist_ok=True)
# 指定最终的索引生成目录,启动索引生成
config = RAGAnythingConfig(
working_dir=WORKING_DIR,
mineru_parse_method="auto",
enable_image_processing=True, # 处理图片
enable_table_processing=True, # 处理表格
enable_equation_processing=True, # 处理公式
)
# 自定义的大模型函数
llm_model_func = create_llm_model_func()
# 自定义的可视模型函数
vision_model_func = create_vision_model_func(llm_model_func)
# 自定义的嵌入函数
embedding_func = create_embedding_func()
rag = RAGAnything(
config=config,
llm_model_func=llm_model_func,
vision_model_func=vision_model_func,
embedding_func=embedding_func,
)
await rag.process_document_complete(
file_path=file_path,
output_dir=output_dir,
parse_method="auto"
)
if __name__ == "__main__":
asyncio.run(main())