Spaces:
Running
Running
File size: 4,859 Bytes
9a1014e | 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 | from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
from typing import Any
import chromadb
from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from pypdf import PdfReader
from hf_text_embeddings import HFTextEmbeddings
DEFAULT_PERSIST_DIR = "./chroma_db_docs"
DEFAULT_COLLECTION = "document_context"
DEFAULT_DOCUMENTS_DIR = "documentos"
PDF_EXTENSIONS = {".pdf"}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Index local documents into Chroma for RAG.")
parser.add_argument("--documents-dir", default=DEFAULT_DOCUMENTS_DIR)
parser.add_argument("--files", help="Comma-separated document paths. Defaults to all PDFs under documents-dir.")
parser.add_argument("--persist-dir", default=DEFAULT_PERSIST_DIR)
parser.add_argument("--collection", default=DEFAULT_COLLECTION)
parser.add_argument("--chunk-size", type=int, default=1200)
parser.add_argument("--chunk-overlap", type=int, default=180)
parser.add_argument("--reset-collection", action="store_true")
return parser.parse_args()
def iter_document_paths(documents_dir: str, files: str | None) -> list[Path]:
if files:
return [Path(item.strip()) for item in files.split(",") if item.strip()]
root = Path(documents_dir)
if not root.exists():
return []
return [
path
for path in sorted(root.rglob("*"))
if path.is_file() and path.suffix.lower() in PDF_EXTENSIONS
]
def clean_pdf_text(text: str) -> str:
lines = [line.strip() for line in text.splitlines()]
return "\n".join(line for line in lines if line)
def metadata_path_for(path: Path) -> Path:
return path.with_suffix(".metadata.json")
def load_document_metadata(path: Path) -> dict[str, Any]:
metadata_path = metadata_path_for(path)
if not metadata_path.exists():
return {}
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
if not isinstance(payload, dict):
raise ValueError(f"Invalid metadata object in {metadata_path}.")
return payload
def load_pdf(path: Path) -> list[Document]:
reader = PdfReader(str(path))
docs: list[Document] = []
document_metadata = load_document_metadata(path)
for page_index, page in enumerate(reader.pages, start=1):
text = clean_pdf_text(page.extract_text() or "")
if not text:
continue
docs.append(
Document(
page_content=text,
metadata={
"source_type": "document",
"source": str(path),
"filename": path.name,
"page": page_index,
"document_title": path.stem,
**document_metadata,
},
)
)
return docs
def stable_id(doc: Document, chunk_index: int) -> str:
source = doc.metadata.get("source", "")
page = doc.metadata.get("page", "")
digest = hashlib.sha1(f"{source}|{page}|{chunk_index}|{doc.page_content}".encode("utf-8")).hexdigest()
return f"document:{Path(str(source)).stem}:p{page}:{digest[:16]}"
def main() -> None:
load_dotenv()
args = parse_args()
paths = iter_document_paths(args.documents_dir, args.files)
if not paths:
raise RuntimeError("No PDF documents found to index.")
if args.reset_collection:
client = chromadb.PersistentClient(path=args.persist_dir)
try:
client.delete_collection(args.collection)
print(f"Deleted existing Chroma collection: {args.collection}")
except Exception:
print(f"Chroma collection did not exist yet: {args.collection}")
embeddings = HFTextEmbeddings()
vectorstore = Chroma(
collection_name=args.collection,
embedding_function=embeddings,
persist_directory=args.persist_dir,
)
splitter = RecursiveCharacterTextSplitter(
chunk_size=args.chunk_size,
chunk_overlap=args.chunk_overlap,
)
total_pages = 0
total_chunks = 0
for path in paths:
docs = load_pdf(path)
splits = splitter.split_documents(docs)
ids = [stable_id(doc, index) for index, doc in enumerate(splits)]
try:
vectorstore.delete(ids=ids)
except Exception:
pass
vectorstore.add_documents(splits, ids=ids)
total_pages += len(docs)
total_chunks += len(splits)
print(f"Indexed {path}: {len(docs)} pages, {len(splits)} chunks")
print(f"Done. Added {total_chunks} chunks from {total_pages} pages into {args.persist_dir} / {args.collection}.")
if __name__ == "__main__":
main()
|