100 lines
3.1 KiB
Python
100 lines
3.1 KiB
Python
import os
|
|
import glob
|
|
import lancedb
|
|
from PyPDF2 import PdfReader
|
|
import openai
|
|
import tkinter as tk
|
|
from tkinter import filedialog
|
|
from config import OPENAI_API_KEY
|
|
|
|
# Setze deinen OpenAI API-Schlüssel
|
|
openai.api_key = OPENAI_API_KEY
|
|
|
|
def extract_text_from_pdf(file_path):
|
|
reader = PdfReader(file_path)
|
|
text = ""
|
|
for page in reader.pages:
|
|
page_text = page.extract_text()
|
|
if page_text:
|
|
text += page_text
|
|
return text
|
|
|
|
def get_embedding(text):
|
|
# OpenAI Embeddings API verwenden
|
|
response = openai.Embedding.create(
|
|
input=text,
|
|
model="text-embedding-ada-002"
|
|
)
|
|
embedding = response['data'][0]['embedding']
|
|
return embedding
|
|
|
|
def load_pdfs_to_lancedb(directory, db_path):
|
|
# Verbindung zur LanceDB herstellen (synchroner Client)
|
|
db = lancedb.connect(db_path)
|
|
table_name = 'pdf_embeddings'
|
|
data = []
|
|
# Über alle PDFs im Verzeichnis iterieren
|
|
for pdf_file in glob.glob(os.path.join(directory, "*.pdf")):
|
|
text = extract_text_from_pdf(pdf_file)
|
|
if not text:
|
|
continue
|
|
embedding = get_embedding(text)
|
|
# Daten zum Einfügen vorbereiten
|
|
data.append({
|
|
'vector': embedding,
|
|
'file_link': pdf_file,
|
|
'text': text # Optional, könnte hilfreich sein
|
|
})
|
|
# Tabelle erstellen oder öffnen
|
|
if table_name in db.table_names():
|
|
table = db.open_table(table_name)
|
|
table.add(data)
|
|
else:
|
|
table = db.create_table(table_name, data=data)
|
|
|
|
def semantic_search(query, db_path, top_k=5):
|
|
db = lancedb.connect(db_path)
|
|
table = db.open_table('pdf_embeddings')
|
|
query_embedding = get_embedding(query)
|
|
# Suche in der Tabelle durchführen
|
|
results = table.search(query_embedding).limit(top_k).to_df()
|
|
# Dateilinks aus den Ergebnissen extrahieren
|
|
file_links = results['file_link'].tolist()
|
|
return file_links
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
import argparse
|
|
|
|
# Argument Parser für Kommandozeilenargumente
|
|
parser = argparse.ArgumentParser(description='PDF-Semantische Suche mit LanceDB')
|
|
parser.add_argument('--db_dir', type=str, help='Pfad zum LanceDB-Verzeichnis', default='lancedb_data')
|
|
args = parser.parse_args()
|
|
|
|
lancedb_path = args.db_dir
|
|
|
|
# Tkinter GUI initialisieren
|
|
root = tk.Tk()
|
|
root.withdraw() # Hauptfenster ausblenden
|
|
|
|
# Ordnerauswahl-Dialog öffnen
|
|
pdf_directory = filedialog.askdirectory(title="Bitte wählen Sie das PDF-Verzeichnis aus")
|
|
|
|
# Überprüfen, ob das Verzeichnis existiert
|
|
if not pdf_directory or not os.path.isdir(pdf_directory):
|
|
print("Kein gültiges Verzeichnis ausgewählt. Das Programm wird beendet.")
|
|
sys.exit(1)
|
|
|
|
# PDFs in LanceDB laden
|
|
load_pdfs_to_lancedb(pdf_directory, lancedb_path)
|
|
|
|
# Semantische Suche durchführen
|
|
while True:
|
|
query = input("Gib deine Suchanfrage ein (oder 'exit' zum Beenden): ")
|
|
if query.lower() == 'exit':
|
|
break
|
|
results = semantic_search(query, lancedb_path)
|
|
print("Passende Dateien:")
|
|
for file_link in results:
|
|
print(file_link)
|