weiweiw
/
Langchain-Chatchat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update DocumentLoader in knowledge_file.py

This commit is contained in:
imClumsyPanda 2023-08-05 21:51:07 +08:00
parent f68e432570
commit 313e590961
1 changed files with 20 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
server/knowledge_base/knowledge_file.py
View File

@ -1,6 +1,20 @@
import os.path import os.path
from server.knowledge_base.utils import (get_file_path) from server.knowledge_base.utils import (get_file_path)
from server.knowledge_base import KnowledgeBase from server.knowledge_base import KnowledgeBase
import sys
LOADER_DICT = {"UnstructuredFileLoader": ['.eml', '.html', '.json', '.md', '.msg', '.rst',
'.rtf', '.txt', '.xml',
'.doc', '.docx', '.epub', '.odt', '.pdf',
'.ppt', '.pptx', '.tsv'], # '.pdf', '.xlsx', '.csv'
"CSVLoader": [".csv"],
}
SUPPORTED_EXTS = [ext for sublist in LOADER_DICT.values() for ext in sublist]
def get_LoaderClass(file_extension):
for LoaderClass, extensions in LOADER_DICT.items():
if file_extension in extensions:
return LoaderClass
class KnowledgeFile: class KnowledgeFile:
@ -12,16 +26,16 @@ class KnowledgeFile:
self.kb = KnowledgeBase.load(knowledge_base_name) self.kb = KnowledgeBase.load(knowledge_base_name)
self.filename = filename self.filename = filename
self.ext = os.path.splitext(filename)[-1] self.ext = os.path.splitext(filename)[-1]
if self.ext not in SUPPORTED_EXTS:
raise ValueError(f"暂未支持的文件格式 {self.ext}")
self.filepath = get_file_path(knowledge_base_name, filename) self.filepath = get_file_path(knowledge_base_name, filename)
self.docs = None self.docs = None
self.loader_class_name = get_LoaderClass(self.ext)
def file2text(self): def file2text(self):
if self.ext in []: LoaderClass = getattr(sys.modules['langchain.document_loaders'], self.loader_class_name)
from langchain.document_loaders import UnstructuredFileLoader loader = LoaderClass(self.filepath)
loader = UnstructuredFileLoader(self.filepath)
elif self.ext in []:
pass
from langchain.text_splitter import CharacterTextSplitter from langchain.text_splitter import CharacterTextSplitter
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=200) text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=200)
return loader.load_and_split(text_splitter) return loader.load_and_split(text_splitter)