Files
Lokales_RAG/loadPDF.py
2024-10-02 16:19:42 +02:00

77 lines
2.5 KiB
Python

import os
import fitz # PyMuPDF zum Extrahieren von PDF-Text
import marqo
# Verbindung zur Marqo-Instanz herstellen
client = marqo.Client(url="http://localhost:8882") # Marqo-URL
# Marqo-Indexname
index_name = "pdf-index"
# Ordner mit PDFs
pdf_folder = "/opt/ki/st01"
# Index-Einstellungen
index_settings = {
"treatUrlsAndPointersAsImages": False,
"model": "hf/e5-base-v2", # Du kannst das Model nach Bedarf anpassen
}
# Funktion zum Extrahieren von Text aus einem PDF
def extract_text_from_pdf(pdf_path):
text = ""
try:
doc = fitz.open(pdf_path)
for page_num in range(doc.page_count):
page = doc.load_page(page_num)
text += page.get_text("text")
doc.close()
except Exception as e:
print(f"Fehler beim Lesen der Datei {pdf_path}: {e}")
return text
# Alle PDFs aus einem Ordner laden und in Marqo importieren
def import_pdfs_to_marqo(folder_path, index_name):
for filename in os.listdir(folder_path):
if filename.endswith(".pdf"):
pdf_path = os.path.join(folder_path, filename)
print(f"Verarbeite {filename}...")
# Text aus PDF extrahieren
pdf_text = extract_text_from_pdf(pdf_path)
if pdf_text:
# Dokument in Marqo importieren
document = {
"pdf_filename": filename,
"pdf_text": pdf_text
}
client.index(index_name).add_documents(
[document],
tensor_fields=["pdf_text"] # Hier wird das Textfeld für die Vektorisierung angegeben
)
print(f"{filename} erfolgreich in Marqo importiert.")
# Funktion zur Überprüfung, ob ein Index existiert
def index_exists(client, index_name):
try:
client.index(index_name).get_settings() # Abfrage, ob der Index existiert
return True
except marqo.errors.IndexNotFoundError:
return False
# Hauptprogramm
if __name__ == "__main__":
# Prüfen, ob der Index existiert
if not index_exists(client, index_name):
print(f"Index {index_name} existiert nicht. Erstelle den Index...")
client.create_index(index_name, settings_dict=index_settings)
print(f"Index {index_name} wurde neu erstellt.")
else:
print(f"Index {index_name} existiert bereits. Importiere neue Dokumente...")
# PDFs importieren
import_pdfs_to_marqo(pdf_folder, index_name)
print("Fertig!")