more file format support

这个提交包含在:
lbykkkk
2024-10-12 18:25:33 +00:00
父节点 d13f1e270c
当前提交 fce4fa1ec7
共有 2 个文件被更改,包括 5 次插入17 次删除

查看文件

@@ -1,22 +1,16 @@
import os
from llama_index.core import SimpleDirectoryReader
# 保留你原有的自定义解析函数
from PyPDF2 import PdfReader
supports_format = ['.csv', '.docx', '.epub', '.ipynb', '.mbox', '.md', '.pdf', '.txt', '.ppt',
'.pptm', '.pptx']
def extract_text_from_pdf(file_path):
reader = PdfReader(file_path)
text = ""
for page in reader.pages:
text += page.extract_text() + "\n"
return text
# 修改后的 extract_text 函数,结合 SimpleDirectoryReader 和自定义解析逻辑
def extract_text(file_path):
_, ext = os.path.splitext(file_path.lower())
# 使用 SimpleDirectoryReader 处理它支持的文件格式
if ext in ['.txt', '.md', '.pdf', '.docx', '.html']:
if ext in supports_format:
try:
reader = SimpleDirectoryReader(input_files=[file_path])
documents = reader.load_data()
@@ -25,10 +19,4 @@ def extract_text(file_path):
except Exception as e:
pass
# 如果 SimpleDirectoryReader 失败,或文件格式不支持,使用自定义解析逻辑
if ext == '.pdf':
try:
return extract_text_from_pdf(file_path)
except Exception as e:
pass
return None