Publish solution for task 6a02e23da6fe2e4ac16acf65: update load_documents.py
This commit is contained in:
+13
-12
@@ -1,14 +1,13 @@
|
|||||||
"""Utility script to load documents from a directory into the Qdrant vector store.
|
# 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
|
# 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
|
# `RecursiveCharacterTextSplitter`, embeds the chunks with `OllamaEmbeddings`, and stores them in the
|
||||||
local Qdrant collection via the helper functions defined in :mod:`main`.
|
# local Chroma collection via the helper functions defined in :mod:`main`.
|
||||||
|
#
|
||||||
Usage:
|
# Usage:
|
||||||
python load_documents.py /path/to/docs
|
# python load_documents.py /path/to/docs
|
||||||
|
#
|
||||||
The script prints the number of documents added.
|
# The script prints the number of documents added.
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
@@ -48,7 +47,9 @@ def load_documents_from_dir(directory: str) -> int:
|
|||||||
title = file_path.stem
|
title = file_path.stem
|
||||||
documents = chunk_document(content, title)
|
documents = chunk_document(content, title)
|
||||||
ids = [str(uuid4()) for _ in documents]
|
ids = [str(uuid4()) for _ in documents]
|
||||||
vector_store.add_documents(documents, ids=ids)
|
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)
|
total_chunks += len(documents)
|
||||||
return total_chunks
|
return total_chunks
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user