cookidooAI/EpubToRag.py

95 lines
3.1 KiB
Python
Executable File

import os
import shutil
import subprocess
import sys
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
# ================= KONFIGURACJA =================
# Podaj tu nazwę swojego pliku EPUB (musi być w tym samym folderze)
EPUB_FILE = "build123d-readthedocs-io-en-latest.epub"
# Nazwa pliku tymczasowego (skrypt sam go stworzy)
MD_FILE = "converted_docs.md"
# Folder bazy danych
DB_DIR = "/home/pali112/db_build123d"
# ================================================
def convert_epub_to_md():
"""Konwertuje EPUB do MD używając systemowego Pandoc."""
print(f"🔄 Rozpoczynam konwersję: {EPUB_FILE} -> {MD_FILE}")
if not os.path.exists(EPUB_FILE):
print(f"❌ BŁĄD: Nie znaleziono pliku {EPUB_FILE}!")
sys.exit(1)
# Sprawdzenie czy pandoc jest zainstalowany
if shutil.which("pandoc") is None:
print("❌ BŁĄD: Nie znaleziono 'pandoc'. Zainstaluj go: sudo apt install pandoc")
sys.exit(1)
try:
# Wywołanie komendy systemowej
cmd = ["pandoc", EPUB_FILE, "-t", "markdown", "-o", MD_FILE]
subprocess.run(cmd, check=True)
print("✅ Konwersja zakończona sukcesem.")
except subprocess.CalledProcessError as e:
print(f"❌ Błąd podczas konwersji pandoc: {e}")
sys.exit(1)
def rebuild_database():
# 1. Konwersja (Automatyczna)
convert_epub_to_md()
# 2. Czyszczenie starej bazy
if os.path.exists(DB_DIR):
print(f"🗑️ Usuwam starą bazę: {DB_DIR}")
shutil.rmtree(DB_DIR)
# 3. Ładowanie przekonwertowanego pliku
print(f"📖 Wczytuję treść z: {MD_FILE}...")
loader = UnstructuredMarkdownLoader(MD_FILE)
data = loader.load()
# Statystyka surowego tekstu
raw_text_len = len(data[0].page_content)
print(f" Wczytano {raw_text_len} znaków tekstu.")
if raw_text_len < 1000:
print("⚠️ UWAGA: Bardzo mało tekstu! Konwersja mogła się nie udać.")
# 4. Dzielenie na kawałki (Chunking)
print("✂️ Dzielę tekst na fragmenty...")
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
separators=["\n## ", "\n### ", "\n", " ", ""]
)
all_splits = text_splitter.split_documents(data)
count = len(all_splits)
print(f" Stworzono {count} chunków.")
# 5. Tworzenie bazy wektorowej
print("🧠 Generowanie embedingów i zapisywanie do ChromaDB...")
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma.from_documents(
documents=all_splits,
embedding=embeddings,
persist_directory=DB_DIR
)
final_count = vectorstore._collection.count()
print("-" * 40)
print(f"✅ SUKCES! Baza gotowa w {DB_DIR}")
print(f"📊 Liczba wektorów w bazie: {final_count}")
print("-" * 40)
# Opcjonalnie: Usuń plik tymczasowy md, żeby nie śmiecić
# os.remove(MD_FILE)
if __name__ == "__main__":
rebuild_database()