RAG_test_1 / app.py
la04's picture
Update app.py
d859c3e verified
raw
history blame
3.21 kB
import faiss
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer
import fitz # PyMuPDF für die Textextraktion aus PDFs
from transformers import pipeline
import logging
# Logging konfigurieren
logging.basicConfig(level=logging.INFO)
# Schritt 1: Lade das Modell für die Embeddings
model = SentenceTransformer('all-mpnet-base-v2') # Besser geeignet für längere Texte
qa_model = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
# FAISS-Index erstellen
def create_faiss_index(documents):
document_embeddings = model.encode(documents)
dimension = len(document_embeddings[0])
index = faiss.IndexFlatL2(dimension)
document_embeddings = np.array(document_embeddings).astype('float32')
index.add(document_embeddings) # Füge Dokumente zum Index hinzu
return index, documents
# Schritt 2: Extrahiere Text aus einem PDF (kleinere Abschnitte, z. B. Absätze)
def extract_text_from_pdf(pdf_path):
doc = fitz.open(pdf_path)
text_chunks = []
for page_num in range(len(doc)):
page = doc.load_page(page_num)
text = page.get_text("text")
# Text in kleinere Abschnitte (z. B. Absätze) unterteilen
chunks = text.split('\n\n') # Unterteilen nach Absätzen
text_chunks.extend(chunks)
return text_chunks
# Schritt 3: Suche nach mehreren passenden Abschnitten
def search_documents(query, index, documents, k=5):
query_embedding = model.encode([query])[0].astype('float32')
D, I = index.search(np.array([query_embedding]), k=k) # Suche nach den Top k Treffern
results = [documents[i] for i in I[0]]
return " ".join(results) # Kombiniere die Top-Ergebnisse
# Schritt 4: Nutze das QA-Modell, um eine präzise Antwort zu generieren
def generate_answer(context, question):
result = qa_model(question=question, context=context)
return result['answer']
# Schritt 5: Gesamtprozess (Fragebeantwortung)
def chatbot_response(pdf_path, question):
logging.info(f"Frage: {question}")
# Textextraktion aus der PDF
text_chunks = extract_text_from_pdf(pdf_path)
# FAISS-Index erstellen
index, documents = create_faiss_index(text_chunks)
# Suche nach Kontext
context = search_documents(question, index, documents, k=5)
logging.info(f"Gefundener Kontext: {context[:500]}") # Nur die ersten 500 Zeichen anzeigen
# Antwort generieren
answer = generate_answer(context, question)
return answer
# Gradio-Interface
pdf_input = gr.File(label="PDF-Datei hochladen", type="filepath") # Ändert type="file" zu type="filepath"
question_input = gr.Textbox(label="Frage eingeben", placeholder="Stelle eine Frage zu dem PDF-Dokument")
response_output = gr.Textbox(label="Antwort")
# Gradio-Interface erstellen
interface = gr.Interface(
fn=chatbot_response,
inputs=[pdf_input, question_input],
outputs=response_output,
title="PDF-Fragebeantwortung mit FAISS und Transformers",
description="Lade eine PDF-Datei hoch und stelle Fragen zu ihrem Inhalt. Das System verwendet FAISS, Transformers und ein QA-Modell, um präzise Antworten zu liefern."
)
if __name__ == "__main__":
interface.launch()