Update app.py

app.py

@@ -1,20 +1,34 @@
 import gradio as gr
 import os
-from langchain.vectorstores import Chroma  # Verwenden des normalen Chroma Moduls
+from langchain.vectorstores import Chroma
+from langchain.document_loaders import PyPDFLoader  # Korrekt importiert
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-from pdfplumber import open as open_pdf  # Verwenden von pdfplumber zum Extrahieren von Text aus PDFs
+from pdf2image import convert_from_path
+from transformers import LayoutLMv3Processor, AutoModelForTokenClassification
 
-def process_pdf_and_query(pdf_path, question):
-    # Lade die PDF und extrahiere den Text
-    with open_pdf(pdf_path) as pdf:
-        text = ""
-        for page in pdf.pages:
-            text += page.extract_text()
+# OCR-Modell einrichten
+class LayoutLMv3OCR:
+    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
+
+ocr_tool = LayoutLMv3OCR()
 
-    # Text als Dokumente in den Chroma Vektor-Datenbank laden
-    documents = [{"content": text, "metadata": {"source": pdf_path}}]
+def process_pdf_and_query(pdf_path, question):
+    loader = PyPDFLoader(pdf_path)
+    documents = loader.load()
 
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     vectordb = Chroma.from_documents(documents, embeddings)
@@ -29,7 +43,7 @@ def process_pdf_and_query(pdf_path, question):
 
 def chatbot_response(pdf, question):
     # Gradio gibt uns die PDF als NamedString, wir extrahieren den Inhalt als Byte-Stream
-    pdf_path = "/tmp/uploaded_pdf.pdf"
+    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
@@ -38,12 +52,12 @@ def chatbot_response(pdf, question):
     with open(pdf_path, "wb") as f:
         f.write(pdf_content)
 
-    # Frage beantworten basierend auf dem Text der PDF
+    # OCR-Text extrahieren
+    extracted_text = ocr_tool.extract_text(pdf_path)
+
+    # Frage beantworten basierend auf der PDF und OCR-Inhalten
     answer = process_pdf_and_query(pdf_path, question)
 
-    # Temporäre Datei löschen
-    os.remove(pdf_path)
-    
     return answer
 
 # Gradio Interface