feat: Chroma + OllamaEmbeddings(nomic-embed-text)
This commit is contained in:
+18
-39
@@ -1,54 +1,33 @@
|
|||||||
import os
|
"""ChromaDB + OllamaEmbeddings (nomic-embed-text)."""
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import List
|
|
||||||
|
|
||||||
from langchain_ollama import OllamaEmbeddings
|
from langchain_ollama import OllamaEmbeddings
|
||||||
from langchain_chroma import Chroma
|
from langchain_chroma import Chroma
|
||||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||||
from langchain.docstore.document import Document
|
from langchain_core.documents import Document
|
||||||
|
|
||||||
|
|
||||||
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
||||||
"""Create or load a Chroma vector store backed by Ollama embeddings.
|
"""Создать/открыть ChromaDB с Ollama-эмбеддингами."""
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
persist_directory: str
|
|
||||||
Directory where Chroma will persist its data.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
Chroma
|
|
||||||
A Chroma vector store instance.
|
|
||||||
"""
|
|
||||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
return Chroma(
|
return Chroma(
|
||||||
persist_directory=persist_directory,
|
collection_name="rag_kb",
|
||||||
embedding_function=embeddings,
|
embedding_function=embeddings,
|
||||||
|
persist_directory=persist_directory,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
|
def load_documents(directory: str, vectorstore: Chroma) -> int:
|
||||||
"""Load all .txt and .md files from *directory* into *vectorstore*.
|
"""Загрузить .txt/.md файлы из директории в ChromaDB с чанкингом."""
|
||||||
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
The function reads files, splits them into chunks using
|
docs = []
|
||||||
``RecursiveCharacterTextSplitter`` and adds the resulting
|
|
||||||
:class:`~langchain.docstore.document.Document` objects to the
|
|
||||||
vector store.
|
|
||||||
"""
|
|
||||||
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
|
||||||
docs: List[Document] = []
|
|
||||||
for path in Path(directory).glob("**/*"):
|
for path in Path(directory).glob("**/*"):
|
||||||
if path.suffix.lower() not in {".txt", ".md"}:
|
if path.suffix.lower() not in (".txt", ".md"):
|
||||||
continue
|
continue
|
||||||
text = path.read_text(encoding="utf-8")
|
text = path.read_text(encoding="utf-8", errors="replace")
|
||||||
docs.extend(splitter.split_text(text))
|
chunks = splitter.split_text(text)
|
||||||
# Convert list of strings to Document objects
|
docs.extend(
|
||||||
documents = [Document(page_content=chunk, metadata={"source": str(p)}) for chunk in docs]
|
[Document(page_content=c, metadata={"source": str(path)}) for c in chunks]
|
||||||
vectorstore.add_documents(documents)
|
)
|
||||||
|
if docs:
|
||||||
# If this module is executed directly, load the default documents folder.
|
vectorstore.add_documents(docs)
|
||||||
if __name__ == "__main__":
|
return len(docs)
|
||||||
vs = create_vectorstore()
|
|
||||||
load_documents("documents", vs)
|
|
||||||
print("Vector store populated.")
|
|
||||||
|
|||||||
Reference in New Issue
Block a user