Spaces:
Sleeping
Sleeping
File size: 3,084 Bytes
17aaccd | 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 | """
Document Loader — ingests PDF/TXT documents from the data/papers directory.
Chunks documents with overlap for better retrieval context.
"""
from __future__ import annotations
import re
from pathlib import Path
def _chunk_text(
text: str,
chunk_size: int = 400,
overlap: int = 80,
) -> list[str]:
"""Split text into overlapping chunks by word count."""
words = text.split()
chunks = []
start = 0
while start < len(words):
end = min(start + chunk_size, len(words))
chunk = " ".join(words[start:end])
if len(chunk.strip()) > 50: # skip trivially short chunks
chunks.append(chunk)
start += chunk_size - overlap
return chunks
def _clean_text(text: str) -> str:
"""Basic text cleanup — remove excessive whitespace and encoding artefacts."""
text = re.sub(r'\s+', ' ', text)
text = re.sub(r'[^\x00-\x7F]+', ' ', text)
return text.strip()
def load_pdf(path: Path) -> list[dict]:
"""Load a PDF and return chunked text with metadata."""
try:
import fitz # PyMuPDF
doc = fitz.open(str(path))
full_text = " ".join(page.get_text() for page in doc)
full_text = _clean_text(full_text)
chunks = _chunk_text(full_text)
meta = {
"source": path.name,
"title": path.stem.replace("_", " ").replace("-", " "),
"authors": "Rahman MA et al.",
"year": "2024",
"domain": "research_paper",
}
return [{"text": c, "metadata": meta} for c in chunks]
except ImportError:
print("PyMuPDF not installed — skipping PDF loading. Run: pip install pymupdf")
return []
except Exception as e:
print(f"Error loading {path}: {e}")
return []
def load_txt(path: Path) -> list[dict]:
"""Load a plain text file and return chunked content."""
try:
text = _clean_text(path.read_text(encoding="utf-8", errors="ignore"))
chunks = _chunk_text(text)
meta = {
"source": path.name,
"title": path.stem.replace("_", " "),
"authors": "",
"year": "",
"domain": "document",
}
return [{"text": c, "metadata": meta} for c in chunks]
except Exception as e:
print(f"Error loading {path}: {e}")
return []
def load_all_documents(docs_dir: Path) -> list:
"""Load all supported documents from a directory."""
from src.retrieval.rag_pipeline import DocumentChunk
docs_dir = Path(docs_dir)
if not docs_dir.exists():
return []
all_chunks = []
loaders = {".pdf": load_pdf, ".txt": load_txt}
for ext, loader in loaders.items():
for fpath in docs_dir.glob(f"*{ext}"):
print(f"Loading: {fpath.name}")
raw = loader(fpath)
for item in raw:
all_chunks.append(
DocumentChunk(text=item["text"], metadata=item["metadata"])
)
print(f"Loaded {len(all_chunks)} chunks from {docs_dir}")
return all_chunks
|