109 lines
3.9 KiB
Python
109 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
init_loader.py
|
||
|
||
Скрипт для загрузки файлов из указанной директории в локальное Qdrant‑хранилище.
|
||
Использует:
|
||
* langchain.document_loaders.FileSystemLoader – чтение всех файлов
|
||
* RecursiveCharacterTextSplitter – разбиение на чанки
|
||
* OllamaEmbeddings (nomic-embed-text) – генерация эмбеддингов
|
||
* QdrantVectorStore – сохранение векторных представлений
|
||
|
||
Параметры:
|
||
--data-dir Путь к директории с исходными документами
|
||
--collection Название коллекции в Qdrant (по умолчанию: rag_collection)
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
# LangChain imports
|
||
from langchain.document_loaders import DirectoryLoader
|
||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||
from langchain.embeddings.ollama import OllamaEmbeddings
|
||
from langchain.vectorstores.qdrant import QdrantVectorStore
|
||
from langchain.schema import Document
|
||
|
||
|
||
def load_documents(data_dir: Path) -> List[Document]:
|
||
"""
|
||
Загружает все файлы из указанной директории в список объектов Document.
|
||
Поддерживаются форматы *.txt, *.md, *.pdf (если установлен pdfminer).
|
||
"""
|
||
loader = DirectoryLoader(
|
||
str(data_dir),
|
||
glob="**/*",
|
||
suffixes=[".txt", ".md", ".pdf"],
|
||
show_progress=True,
|
||
)
|
||
return loader.load()
|
||
|
||
|
||
def chunk_documents(docs: List[Document], chunk_size: int = 1000, overlap: int = 200) -> List[Document]:
|
||
"""
|
||
Разбивает каждый документ на чанки фиксированного размера.
|
||
"""
|
||
splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=chunk_size,
|
||
chunk_overlap=overlap,
|
||
separators=["\n\n", "\n", " ", ""],
|
||
)
|
||
return splitter.split_documents(docs)
|
||
|
||
|
||
def embed_and_store(chunks: List[Document], collection_name: str) -> None:
|
||
"""
|
||
Генерирует эмбеддинги для чанков и сохраняет их в Qdrant.
|
||
"""
|
||
# Инициализируем Ollama embeddings
|
||
embedding = OllamaEmbeddings(model="nomic-embed-text")
|
||
|
||
# Создаём (или подключаемся к) коллекцию Qdrant
|
||
vector_store = QdrantVectorStore(
|
||
collection_name=collection_name,
|
||
embedding=embedding,
|
||
url="http://localhost:6333",
|
||
)
|
||
|
||
# Добавляем документы в хранилище
|
||
vector_store.add_documents(chunks)
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="Загрузка документов в Qdrant.")
|
||
parser.add_argument(
|
||
"--data-dir",
|
||
type=str,
|
||
required=True,
|
||
help="Путь к директории с документами.",
|
||
)
|
||
parser.add_argument(
|
||
"--collection",
|
||
type=str,
|
||
default="rag_collection",
|
||
help="Имя коллекции в Qdrant.",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
data_dir = Path(args.data_dir).expanduser().resolve()
|
||
if not data_dir.is_dir():
|
||
raise FileNotFoundError(f"Каталог {data_dir} не найден.")
|
||
|
||
print(f"[INFO] Загружаем документы из: {data_dir}")
|
||
docs = load_documents(data_dir)
|
||
print(f"[INFO] Найдено {len(docs)} документов. Разбиваем на чанки...")
|
||
|
||
chunks = chunk_documents(docs)
|
||
print(f"[INFO] Получено {len(chunks)} чанков. Генерируем эмбеддинги и сохраняем в Qdrant...")
|
||
|
||
embed_and_store(chunks, args.collection)
|
||
|
||
print("[SUCCESS] Загрузка завершена.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |