File size: 1,454 Bytes
cdc87cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """
Loads TXT and PDF files from the /data directory.
"""
import os
from langchain_community.document_loaders import TextLoader, PyPDFLoader
def load_documents(data_dir: str) -> list:
"""
Load all .txt and .pdf documents from the given directory.
Returns a list of LangChain Document objects.
"""
documents = []
supported_extensions = (".txt", ".pdf")
for filename in os.listdir(data_dir):
if not filename.lower().endswith(supported_extensions):
continue
filepath = os.path.join(data_dir, filename)
try:
if filename.lower().endswith(".pdf"):
loader = PyPDFLoader(filepath)
else:
loader = TextLoader(filepath, encoding="utf-8")
docs = loader.load()
# Tag each doc with source filename
for doc in docs:
doc.metadata["source"] = filename
documents.extend(docs)
print(f" [OK] Loaded: {filename} ({len(docs)} page(s))")
except Exception as e:
print(f" [WARN] Could not load {filename}: {e}")
print(f"\n Total documents loaded: {len(documents)}")
return documents
if __name__ == "__main__":
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
data_folder = os.path.join(base_dir, "data")
docs = load_documents(data_folder)
print(f"Sample content preview:\n{docs[0].page_content[:300]}")
|