27 lines
1.1 KiB
Python
27 lines
1.1 KiB
Python
import os
|
|
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"):
|
|
if not os.path.exists(persist_directory):
|
|
os.makedirs(persist_directory, exist_ok=True)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
def load_documents(directory: str, vectorstore) -> None:
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
docs: List[Document] = []
|
|
for path in Path(directory).rglob("*.txt"):
|
|
text = path.read_text(encoding="utf-8")
|
|
docs.extend(splitter.split_documents([Document(page_content=text)]))
|
|
for path in Path(directory).rglob("*.md"):
|
|
text = path.read_text(encoding="utf-8")
|
|
docs.extend(splitter.split_documents([Document(page_content=text)]))
|
|
if docs:
|
|
vectorstore.add_documents(docs)
|