Update app.py

app.py

@@ -1,66 +1,93 @@
 import gradio as gr
-import os
-from langchain.vectorstores import Chroma  # Verwende das Standard-Chroma-Modul
-from langchain.document_loaders import PyPDFLoader  # Korrekt importiert
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from transformers import LayoutLMv3Processor, AutoModelForTokenClassification
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-from pdf2image import convert_from_path
-from transformers import LayoutLMv3Processor, AutoModelForTokenClassification
+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-Modell einrichten
-class LayoutLMv3OCR:
+# 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):
-        images = convert_from_path(pdf_path)
-        text_pages = []
-        for image in images:
-            inputs = self.processor(images=image, return_tensors="pt")
-            outputs = self.model(**inputs)
-            text = self.processor.batch_decode(outputs.logits, skip_special_tokens=True)[0]
-            text_pages.append(text)
-        return text_pages
+        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_tool = LayoutLMv3OCR()
+# OCR-Instanz erstellen
+ocr_tool = OCRTool()
 
+# Funktion zur Verarbeitung der PDF und Antwortgenerierung
 def process_pdf_and_query(pdf_path, question):
-    loader = PyPDFLoader(pdf_path)
-    documents = loader.load()
+    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)
+        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)
+        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)
 
-    qa_chain = RetrievalQA.from_chain_type(llm=None, retriever=retriever, chain_type_kwargs={"prompt": prompt})
-    response = qa_chain.run(input_documents=documents, question=question)
-    return response
+        # 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}
+        )
 
-def chatbot_response(pdf, question):
-    # Gradio gibt uns die PDF als NamedString, wir extrahieren den Inhalt als Byte-Stream
-    pdf_path = "/mnt/data/uploaded_pdf.pdf"  # Verwendet das persistente Verzeichnis von Hugging Face Spaces
-    
-    # Extrahiere den Inhalt der Datei als Bytes
-    pdf_content = pdf.read()  # Hier holen wir den Inhalt der PDF als Byte-Stream
-
-    # Speichern des Byte-Streams von der Datei
-    with open(pdf_path, "wb") as f:
-        f.write(pdf_content)
+        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."
 
-    # OCR-Text extrahieren
-    extracted_text = ocr_tool.extract_text(pdf_path)
+# 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)
 
-    # Frage beantworten basierend auf der PDF und OCR-Inhalten
-    answer = process_pdf_and_query(pdf_path, question)
+        # 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."
 
-    return answer
+        # 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
+# Gradio-Interface
 pdf_input = gr.File(label="PDF-Datei hochladen")
 question_input = gr.Textbox(label="Frage eingeben")
 response_output = gr.Textbox(label="Antwort")