115 lines
4.5 KiB
Python
115 lines
4.5 KiB
Python
import json
|
||
import os
|
||
import sys
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import torch
|
||
from langchain_chroma import Chroma
|
||
from langchain_huggingface import HuggingFaceEmbeddings
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from loguru import logger
|
||
|
||
# Добавляем путь к родительской директории для импорта config
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
from config import settings
|
||
|
||
|
||
def load_json_files(directory: str) -> List[Dict[str, Any]]:
|
||
"""Загрузка всех JSON файлов из указанной директории."""
|
||
documents = []
|
||
|
||
try:
|
||
if not os.path.exists(directory):
|
||
logger.error(f"Директория {directory} не существует")
|
||
return documents
|
||
|
||
for filename in os.listdir(directory):
|
||
if filename.endswith(".json"):
|
||
file_path = os.path.join(directory, filename)
|
||
try:
|
||
with open(file_path, "r", encoding="utf-8") as file:
|
||
data = json.load(file)
|
||
documents.append(
|
||
{"text": data["text"], "metadata": data["metadata"]}
|
||
)
|
||
logger.info(f"Загружен файл: {filename}")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при чтении файла {filename}: {e}")
|
||
|
||
logger.success(f"Загружено {len(documents)} JSON файлов")
|
||
return documents
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при загрузке JSON файлов: {e}")
|
||
return documents
|
||
|
||
|
||
def split_text_into_chunks(text: str, metadata: Dict[str, Any]) -> List[Any]:
|
||
"""Разделение текста на чанки с сохранением метаданных."""
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=settings.MAX_CHUNK_SIZE,
|
||
chunk_overlap=settings.CHUNK_OVERLAP,
|
||
length_function=len,
|
||
is_separator_regex=False,
|
||
)
|
||
|
||
chunks = text_splitter.create_documents(texts=[text], metadatas=[metadata])
|
||
return chunks
|
||
|
||
|
||
def generate_chroma_db() -> Optional[Chroma]:
|
||
"""Инициализация ChromaDB с данными из JSON файлов."""
|
||
try:
|
||
# Создаем директорию для хранения базы данных, если она не существует
|
||
os.makedirs(settings.DOCS_CHROMA_PATH, exist_ok=True)
|
||
|
||
# Загружаем JSON файлы
|
||
documents = load_json_files(settings.PARSED_JSON_PATH)
|
||
|
||
if not documents:
|
||
logger.warning("Нет документов для добавления в базу данных")
|
||
return None
|
||
|
||
# Инициализируем модель эмбеддингов
|
||
embeddings = HuggingFaceEmbeddings(
|
||
model_name=settings.LM_MODEL_NAME,
|
||
model_kwargs={"device": "cuda" if torch.cuda.is_available() else "cpu"},
|
||
encode_kwargs={"normalize_embeddings": True},
|
||
)
|
||
|
||
# Подготавливаем данные для Chroma
|
||
all_chunks = []
|
||
for i, doc in enumerate(documents):
|
||
chunks = split_text_into_chunks(doc["text"], doc["metadata"])
|
||
all_chunks.extend(chunks)
|
||
logger.info(
|
||
f"Документ {i+1}/{len(documents)} разбит на {len(chunks)} чанков"
|
||
)
|
||
|
||
# Создаем векторное хранилище
|
||
texts = [chunk.page_content for chunk in all_chunks]
|
||
metadatas = [chunk.metadata for chunk in all_chunks]
|
||
ids = [f"doc_{i}" for i in range(len(all_chunks))]
|
||
|
||
chroma_db = Chroma.from_texts(
|
||
texts=texts,
|
||
embedding=embeddings,
|
||
ids=ids,
|
||
metadatas=metadatas,
|
||
persist_directory=settings.DOCS_CHROMA_PATH,
|
||
collection_name=settings.DOCS_COLLECTION_NAME,
|
||
collection_metadata={
|
||
"hnsw:space": "cosine",
|
||
},
|
||
)
|
||
|
||
logger.success(
|
||
f"База Chroma инициализирована, добавлено {len(all_chunks)} чанков из {len(documents)} документов"
|
||
)
|
||
return chroma_db
|
||
except Exception as e:
|
||
logger.error(f"Ошибка инициализации Chroma: {e}")
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
generate_chroma_db() |