54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
"""Module for creating and loading a Chroma vector store with Ollama embeddings."""
|
|
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_chroma import Chroma
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain.docstore.document import Document
|
|
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Create a Chroma vector store with Ollama embeddings.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the vector store will be persisted.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
The created Chroma vector store.
|
|
"""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
|
|
"""Load .txt and .md files from a directory, split them into chunks, and add to the vector store.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Path to the directory containing documents.
|
|
vectorstore: Chroma
|
|
The vector store to add documents to.
|
|
chunk_size: int
|
|
Maximum size of each chunk.
|
|
chunk_overlap: int
|
|
Number of characters to overlap between chunks.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
|
documents: List[Document] = []
|
|
for file_path in Path(directory).glob("**/*"):
|
|
if file_path.suffix.lower() in {".txt", ".md"}:
|
|
text = file_path.read_text(encoding="utf-8")
|
|
chunks = splitter.split_text(text)
|
|
for chunk in chunks:
|
|
documents.append(Document(page_content=chunk, metadata={"source": str(file_path)}))
|
|
if documents:
|
|
vectorstore.add_documents(documents)
|
|
vectorstore.persist()
|