add vectorstore
This commit is contained in:
@@ -0,0 +1,22 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user