Add init_loader.py
This commit is contained in:
@@ -0,0 +1,32 @@
|
|||||||
|
"""Script to load documents from a directory into the vector store."""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from .vector_store import VectorStore
|
||||||
|
from .splitter import chunk_text
|
||||||
|
|
||||||
|
|
||||||
|
def load_documents(directory: str, store: VectorStore | None = None) -> int:
|
||||||
|
"""Load all text files from *directory* into the vector store.
|
||||||
|
|
||||||
|
Returns the number of documents added.
|
||||||
|
"""
|
||||||
|
if store is None:
|
||||||
|
store = VectorStore()
|
||||||
|
count = 0
|
||||||
|
for path in Path(directory).rglob("*.txt"):
|
||||||
|
content = path.read_text(encoding="utf-8")
|
||||||
|
chunks = chunk_text(content)
|
||||||
|
store.add_documents(chunks)
|
||||||
|
count += 1
|
||||||
|
return count
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(description="Load documents into the RAG vector store.")
|
||||||
|
parser.add_argument("directory", help="Path to directory containing .txt files")
|
||||||
|
args = parser.parse_args()
|
||||||
|
added = load_documents(args.directory)
|
||||||
|
print(f"Loaded {added} documents into the vector store.")
|
||||||
Reference in New Issue
Block a user