120 lines
4.8 KiB
Python
120 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
init_loader.py
|
||
|
||
Скрипт для загрузки файлов из указанной директории в локальное векторное хранилище Qdrant.
|
||
Использует:
|
||
- LangChain OllamaEmbeddings (model "nomic-embed-text")
|
||
- RecursiveCharacterTextSplitter для разбиения на чанки
|
||
- langchain_qdrant.QdrantVectorStore для сохранения эмбеддингов
|
||
|
||
Параметры запуска:
|
||
python init_loader.py <директория_с_файлами> [--collection <имя_коллекции>]
|
||
|
||
Пример:
|
||
python init_loader.py ./docs --collection knowledge_base
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
# LangChain imports
|
||
from langchain.embeddings.ollama import OllamaEmbeddings
|
||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||
from langchain.vectorstores.qdrant import QdrantVectorStore
|
||
from langchain.document_loaders import TextLoader, UnstructuredFileLoader
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Конфигурация по умолчанию
|
||
DEFAULT_COLLECTION_NAME = "rag_collection"
|
||
QDRANT_HOST = "localhost"
|
||
QDRANT_PORT = 6333
|
||
EMBEDDING_MODEL = "nomic-embed-text"
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
def _get_text_loaders(directory: Path) -> Iterable[TextLoader]:
|
||
"""
|
||
Возвращает генератор загрузчиков для всех файлов в директории.
|
||
Поддерживаются текстовые файлы и любые форматы, которые умеет UnstructuredFileLoader.
|
||
"""
|
||
for file_path in directory.rglob("*"):
|
||
if file_path.is_file():
|
||
# Для простоты используем TextLoader для .txt,
|
||
# а остальные обрабатываем через UnstructuredFileLoader
|
||
if file_path.suffix.lower() == ".txt":
|
||
yield TextLoader(str(file_path), encoding="utf-8")
|
||
else:
|
||
yield UnstructuredFileLoader(str(file_path))
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
def load_directory_to_qdrant(
|
||
directory: Path,
|
||
collection_name: str = DEFAULT_COLLECTION_NAME,
|
||
) -> None:
|
||
"""
|
||
Загружает все документы из указанной директории в Qdrant.
|
||
|
||
Args:
|
||
directory (Path): Путь к каталогу с документами.
|
||
collection_name (str): Имя коллекции в Qdrant.
|
||
"""
|
||
if not directory.is_dir():
|
||
raise ValueError(f"Путь {directory} не является директорией")
|
||
|
||
# 1. Инициализируем эмбеддер
|
||
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||
|
||
# 2. Создаём текстовый splitter
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=1000,
|
||
chunk_overlap=200,
|
||
)
|
||
|
||
# 3. Инициализируем QdrantVectorStore (создаст коллекцию при необходимости)
|
||
vector_store = QdrantVectorStore.from_existing_collection(
|
||
embedding=embeddings,
|
||
collection_name=collection_name,
|
||
url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
|
||
)
|
||
|
||
# 4. Загружаем и разбиваем документы, затем сохраняем
|
||
for loader in _get_text_loaders(directory):
|
||
documents = loader.load()
|
||
if not documents:
|
||
continue
|
||
|
||
# Разбиваем каждый документ на чанки
|
||
chunks = text_splitter.split_documents(documents)
|
||
|
||
# Добавляем векторные представления в Qdrant
|
||
vector_store.add_documents(chunks)
|
||
|
||
print(f"Загрузка завершена. Коллекция '{collection_name}' содержит "
|
||
f"{vector_store._client.count(collection_name=collection_name).count} точек.")
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
def _parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description="Загрузить документы из директории в Qdrant."
|
||
)
|
||
parser.add_argument(
|
||
"directory",
|
||
type=str,
|
||
help="Путь к каталогу с документами.",
|
||
)
|
||
parser.add_argument(
|
||
"--collection",
|
||
type=str,
|
||
default=DEFAULT_COLLECTION_NAME,
|
||
help=f"Имя коллекции (по умолчанию: {DEFAULT_COLLECTION_NAME}).",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
if __name__ == "__main__":
|
||
args = _parse_args()
|
||
load_directory_to_qdrant(Path(args.directory), collection_name=args.collection) |