langchain本身的DirectoryLoader
仅支持单种文件格式的上传。但是需求是给定一个文件夹,上传所有的文件。所以写了一个FileLoader
以下代码供参考~
DirectoryLoader
本身是基于其他Loader比如 CSVLoader
写的,关于load的参数需要以字典格式传入。
class FileLoader():
def __init__(self, directory_path):
# Define a dictionary to map file extensions to their respective loaders
self.directory_path = directory_path
self.loaders = {
'.pdf': PyMuPDFLoader,
'.csv': CSVLoader,
'.json': JSONLoader,
'.html': BSHTMLLoader,
'.txt': TextLoader
}
# Define a function to create a DirectoryLoader for a specific file type
def create_directory_loader(self, file_type, loader_kwargs: Optional[dict] = None):
return DirectoryLoader(
path=self.directory_path,
glob=f"**/*{file_type}",
loader_cls=self.loaders[file_type],
loader_kwargs=loader_kwargs,
)
def load_all(self):
# Create DirectoryLoader instances for each file type
pdf_loader = self.create_directory_loader('.pdf')
csv_loader = self.create_directory_loader('.csv', {'encoding': 'utf-8'})
html_loader = self.create_directory_loader('.html',
{'open_encoding': 'utf-8', 'bs_kwargs': {'features': 'html.parser'}})
json_loader = self.create_directory_loader('.json', )
txt_loader = self.create_directory_loader('.txt', {'encoding': 'utf-8'})
# Load the files
pdf_documents = pdf_loader.load()
html_documents = html_loader.load()
csv_documents = csv_loader.load()
json_documents = json_loader.load()
txt_documents = txt_loader.load()
return pdf_documents + html_documents + csv_documents + json_documents + txt_documents