68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
# Utility script to load documents from a directory into the Chroma vector store.
|
|
#
|
|
# The script walks through the specified directory, reads all .txt files, splits them into chunks using
|
|
# `RecursiveCharacterTextSplitter`, embeds the chunks with `OllamaEmbeddings`, and stores them in the
|
|
# local Chroma collection via the helper functions defined in :mod:`main`.
|
|
#
|
|
# Usage:
|
|
# python load_documents.py /path/to/docs
|
|
#
|
|
# The script prints the number of documents added.
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from main import get_vector_store, get_embeddings, chunk_document
|
|
|
|
|
|
def load_documents_from_dir(directory: str) -> int:
|
|
"""Load all .txt files from *directory* into the vector store.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Path to the directory containing text files.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
Number of documents (chunks) added to the store.
|
|
"""
|
|
directory_path = Path(directory)
|
|
if not directory_path.is_dir():
|
|
raise ValueError(f"{directory!r} is not a directory")
|
|
|
|
vector_store = get_vector_store()
|
|
embedding = get_embeddings()
|
|
|
|
total_chunks = 0
|
|
for file_path in directory_path.rglob("*.txt"):
|
|
try:
|
|
content = file_path.read_text(encoding="utf-8")
|
|
except Exception as exc:
|
|
print(f"Failed to read {file_path}: {exc}", file=sys.stderr)
|
|
continue
|
|
title = file_path.stem
|
|
documents = chunk_document(content, title)
|
|
ids = [str(uuid4()) for _ in documents]
|
|
embeddings = embedding.embed_documents([doc.page_content for doc in documents])
|
|
collection = vector_store.get_collection(name="rag_memory")
|
|
collection.add(ids=ids, documents=[doc.page_content for doc in documents], embeddings=embeddings, metadatas=[doc.metadata for doc in documents])
|
|
total_chunks += len(documents)
|
|
return total_chunks
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 2:
|
|
print("Usage: python load_documents.py <directory>")
|
|
sys.exit(1)
|
|
dir_path = sys.argv[1]
|
|
try:
|
|
added = load_documents_from_dir(dir_path)
|
|
print(f"Added {added} chunk(s) to the knowledge base.")
|
|
except Exception as exc:
|
|
print(f"Error: {exc}", file=sys.stderr)
|
|
sys.exit(1)
|