| import mimetypes |
| import os |
|
|
| from langchain_community.document_loaders import PyPDFLoader, TextLoader |
|
|
|
|
| def load_documents(data_dir): |
| """ |
| Load supported documents from a file path or directory. |
| """ |
|
|
| documents = [] |
| data_dir = str(data_dir) |
|
|
| if data_dir.endswith(".txt") or data_dir.endswith(".pdf"): |
| loader = _build_loader(data_dir) |
| return loader.load() if loader else [] |
|
|
| if os.path.isdir(data_dir): |
| for root, _, files in os.walk(data_dir): |
| for file_name in files: |
| file_path = os.path.join(root, file_name) |
| loader = _build_loader(file_path) |
| if loader is None: |
| continue |
| documents.extend(loader.load()) |
| return documents |
|
|
| print(f"Invalid path: {data_dir}") |
| return [] |
|
|
|
|
| def _build_loader(file_path): |
| data_type = mimetypes.guess_type(file_path)[0] |
| if data_type == "text/plain": |
| return TextLoader(file_path) |
| if data_type == "application/pdf": |
| return PyPDFLoader(file_path) |
|
|
| print(f"Unsupported file type: {file_path}") |
| return None |
|
|
