Spaces:
Running
Running
File size: 5,303 Bytes
a99c2d3 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """
document_loader.py
------------------
Handles loading and extracting text from different file types.
Supported formats:
- .txt (plain text)
- .pdf (PDF documents)
- .csv (comma-separated values)
- .docx (Microsoft Word documents)
Each loader returns a list of LangChain Document objects.
A Document has two fields:
- page_content : the extracted text
- metadata : a dict with extra info like the source file name
"""
import os
from langchain_core.documents import Document
# ── helpers ──────────────────────────────────────────────────────────────────
def _make_doc(text: str, source: str) -> Document:
"""Wrap extracted text in a LangChain Document with source metadata."""
return Document(page_content=text, metadata={"source": source})
# ── per-format loaders ────────────────────────────────────────────────────────
def load_txt(file_path: str) -> list[Document]:
"""Load a plain-text file and return it as a single Document."""
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
text = f.read()
return [_make_doc(text, file_path)]
def load_pdf(file_path: str) -> list[Document]:
"""
Load a PDF file page-by-page.
Each page becomes its own Document so we can cite the exact page later.
Requires: pypdf
"""
try:
from pypdf import PdfReader
except ImportError:
raise ImportError("pypdf is required for PDF support. Run: pip install pypdf")
reader = PdfReader(file_path)
documents = []
for page_num, page in enumerate(reader.pages):
text = page.extract_text() or ""
if text.strip(): # skip blank pages
doc = Document(
page_content=text,
metadata={"source": file_path, "page": page_num + 1},
)
documents.append(doc)
return documents
def load_csv(file_path: str) -> list[Document]:
"""
Load a CSV file.
Each row is turned into a readable 'key: value' string and stored as
one Document so every row is individually searchable.
Requires: pandas
"""
try:
import pandas as pd
except ImportError:
raise ImportError("pandas is required for CSV support. Run: pip install pandas")
df = pd.read_csv(file_path)
documents = []
for idx, row in df.iterrows():
# Build a human-readable string from each row
row_text = "\n".join(f"{col}: {val}" for col, val in row.items())
doc = Document(
page_content=row_text,
metadata={"source": file_path, "row": idx + 1},
)
documents.append(doc)
return documents
def load_docx(file_path: str) -> list[Document]:
"""
Load a Microsoft Word (.docx) file.
Each paragraph becomes its own Document.
Requires: python-docx
"""
try:
from docx import Document as WordDocument
except ImportError:
raise ImportError(
"python-docx is required for DOCX support. Run: pip install python-docx"
)
word_doc = WordDocument(file_path)
documents = []
for para_num, para in enumerate(word_doc.paragraphs):
text = para.text.strip()
if text: # skip empty paragraphs
doc = Document(
page_content=text,
metadata={"source": file_path, "paragraph": para_num + 1},
)
documents.append(doc)
return documents
# ── main entry point ──────────────────────────────────────────────────────────
def load_document(file_path: str) -> list[Document]:
"""
Detect the file extension and call the right loader.
Parameters
----------
file_path : str
Full path to the file on disk.
Returns
-------
list[Document]
A list of LangChain Document objects with extracted text.
Raises
------
ValueError – if the file type is not supported.
Exception – if loading fails for any reason.
"""
if not os.path.exists(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
extension = os.path.splitext(file_path)[1].lower()
loaders = {
".txt": load_txt,
".pdf": load_pdf,
".csv": load_csv,
".docx": load_docx,
}
if extension not in loaders:
raise ValueError(
f"Unsupported file type: '{extension}'. "
f"Supported types: {', '.join(loaders.keys())}"
)
# Call the appropriate loader
documents = loaders[extension](file_path)
if not documents:
raise ValueError(f"No readable text found in: {file_path}")
print(f" OK: Loaded {len(documents)} chunk(s) from '{os.path.basename(file_path)}'")
return documents
|