56 lines
2.2 KiB
Python
56 lines
2.2 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") -> Chroma:
|
|
"""
|
|
Создаёт и возвращает объект Chroma, подключённый к указанной директории.
|
|
Если директория не существует, она будет создана.
|
|
"""
|
|
Path(persist_directory).mkdir(parents=True, exist_ok=True)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(
|
|
persist_directory=persist_directory,
|
|
embedding_function=embeddings,
|
|
)
|
|
|
|
|
|
def _load_text_files(directory: str) -> List[Document]:
|
|
"""
|
|
Читает все .txt и .md файлы из указанной директории и возвращает список Document.
|
|
"""
|
|
docs: List[Document] = []
|
|
for file_path in Path(directory).rglob("*"):
|
|
if file_path.suffix.lower() in {".txt", ".md"}:
|
|
text = file_path.read_text(encoding="utf-8")
|
|
docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
|
|
return docs
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
|
"""
|
|
Загружает документы из указанной директории в ChromaDB.
|
|
Делает чанкинг с помощью RecursiveCharacterTextSplitter и добавляет в коллекцию.
|
|
"""
|
|
# Читаем файлы
|
|
raw_docs = _load_text_files(directory)
|
|
|
|
# Разбиваем на чанки
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=200,
|
|
separators=["\n\n", "\n", " ", ""],
|
|
)
|
|
chunks = splitter.split_documents(raw_docs)
|
|
|
|
# Добавляем в коллекцию
|
|
vectorstore.add_documents(chunks)
|
|
# Сохраняем изменения
|
|
vectorstore.persist()
|
|
print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}") |