23 lines
982 B
Python
23 lines
982 B
Python
import os
|
|
from pathlib import Path
|
|
from langchain.vectorstores import Chroma
|
|
from langchain.embeddings.ollama import OllamaEmbeddings
|
|
from langchain.document_loaders import TextLoader, MarkdownLoader
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db"):
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(embedding_function=embeddings, persist_directory=persist_directory)
|
|
|
|
def load_documents(directory: str, vectorstore):
|
|
docs = []
|
|
for file in Path(directory).glob("**/*.*"):
|
|
if file.suffix.lower() not in {".txt", ".md"}:
|
|
continue
|
|
loader = TextLoader(str(file)) if file.suffix == ".txt" else MarkdownLoader(str(file))
|
|
docs.extend(loader.load())
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
texts = splitter.split_documents(docs)
|
|
vectorstore.add_documents(texts)
|
|
vectorstore.persist()
|