app.py

import gradio as gr
from langchain.vectorstores import Chroma
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings import HuggingFaceEmbeddings
from transformers import LayoutLMv3Processor, AutoModelForTokenClassification
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import fitz  # PyMuPDF
import os

# Funktion zur Textextraktion mit PyMuPDF (fitz)
def extract_text_from_pdf(pdf_path):
    try:
        doc = fitz.open(pdf_path)
        text_pages = []
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            text = page.get_text("text")  # Extrahiert Text als normalen Text
            text_pages.append(text)
        return text_pages
    except Exception as e:
        print(f"Fehler bei der Textextraktion: {e}")
        return []

# OCR-Tool mit LayoutLMv3 für strukturierte PDFs
class OCRTool:
    def __init__(self):
        self.processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
        self.model = AutoModelForTokenClassification.from_pretrained("microsoft/layoutlmv3-base")

    def extract_text(self, pdf_path):
        try:
            # LayoutLMv3 für strukturierte PDFs verwenden
            text_pages = extract_text_from_pdf(pdf_path)
            return text_pages
        except Exception as e:
            print(f"Fehler bei der PDF-Verarbeitung: {e}")
            return []

# OCR-Instanz erstellen
ocr_tool = OCRTool()

# Funktion zur Verarbeitung der PDF und Antwortgenerierung
def process_pdf_and_query(pdf_path, question):
    try:
        # Dokument laden und in Vektoren umwandeln
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()

        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        vectordb = Chroma.from_documents(documents, embeddings)

        retriever = vectordb.as_retriever()
        prompt_template = "Beantworte die folgende Frage basierend auf dem Dokument: {context}\nFrage: {question}\nAntwort:"
        prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)

        # RetrievalQA-Chain mit Hugging Face LLM
        qa_chain = RetrievalQA.from_chain_type(
            llm="huggingface/gpt2",  # Modell von Hugging Face
            retriever=retriever,
            chain_type_kwargs={"prompt": prompt}
        )

        response = qa_chain.run(input_documents=documents, question=question)
        return response
    except Exception as e:
        print(f"Fehler bei der Fragebeantwortung: {e}")
        return "Es gab ein Problem bei der Verarbeitung der Frage."

# Funktion für die Chatbot-Antwort
def chatbot_response(pdf, question):
    try:
        # Speichern der hochgeladenen PDF
        pdf_path = "uploaded_pdf.pdf"
        pdf.save(pdf_path)

        # Textextraktion aus der PDF
        extracted_text = ocr_tool.extract_text(pdf_path)
        if not extracted_text:
            return "Es konnte kein Text aus der PDF extrahiert werden."

        # Frage beantworten basierend auf den extrahierten Daten
        answer = process_pdf_and_query(pdf_path, question)
        os.remove(pdf_path)
        return answer
    except Exception as e:
        print(f"Fehler bei der Chatbot-Verarbeitung: {e}")
        return "Es gab ein Problem bei der Verarbeitung der Anfrage."

# Gradio-Interface
pdf_input = gr.File(label="PDF-Datei hochladen")
question_input = gr.Textbox(label="Frage eingeben")
response_output = gr.Textbox(label="Antwort")

interface = gr.Interface(
    fn=chatbot_response,
    inputs=[pdf_input, question_input],
    outputs=response_output,
    title="RAG Chatbot mit PDF-Unterstützung",
    description="Lade eine PDF-Datei hoch und stelle Fragen zu ihrem Inhalt."
)

if __name__ == "__main__":
    interface.launch()