Spaces:
Running on Zero
Running on Zero
File size: 2,618 Bytes
9936912 9009a09 9936912 d08774c 9936912 9009a09 9936912 | 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 | """Document loader supporting PDF, Markdown, Text, and JSONL corpora."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
class Document:
def __init__(self, content: str, source_path: str, metadata: dict[str, Any] | None = None) -> None:
self.content = content
self.source_path = source_path
self.metadata = metadata or {}
def to_dict(self) -> dict[str, Any]:
return {
"content": self.content,
"source_path": self.source_path,
"metadata": self.metadata,
}
from controlai_rag.textfix import repair
def load_pdf(path: Path) -> list[Document]:
docs = []
try:
from pypdf import PdfReader
reader = PdfReader(str(path))
for idx, page in enumerate(reader.pages, 1):
# Repair broken symbol-font extraction before the text is ever
# chunked or indexed -- see controlai_rag/textfix.py.
text = repair(page.extract_text() or "")
if text.strip():
docs.append(Document(
content=text.strip(),
source_path=str(path),
metadata={"page": idx, "filename": path.name, "doc_type": "pdf"},
))
except Exception as exc:
print(f"Warning: Failed to load PDF {path}: {exc}")
return docs
def load_text(path: Path) -> list[Document]:
try:
text = path.read_text(encoding="utf-8", errors="ignore")
if text.strip():
return [Document(
content=text.strip(),
source_path=str(path),
metadata={"filename": path.name, "doc_type": path.suffix.lstrip(".")},
)]
except Exception as exc:
print(f"Warning: Failed to load text file {path}: {exc}")
return []
def load_single_file(path: Path) -> list[Document]:
"""Load a single file based on its extension."""
ext = path.suffix.lower()
if ext == ".pdf":
return load_pdf(path)
elif ext in (".md", ".txt"):
return load_text(path)
return []
def load_directory(dir_path: Path) -> list[Document]:
docs = []
if not dir_path.exists():
return docs
for file_path in dir_path.rglob("*"):
if not file_path.is_file() or file_path.name.startswith("."):
continue
suffix = file_path.suffix.lower()
if suffix == ".pdf":
docs.extend(load_pdf(file_path))
elif suffix in (".md", ".txt", ".rst", ".py", ".m", ".json"):
docs.extend(load_text(file_path))
return docs
|