Update app.py

app.py

@@ -1,55 +1,64 @@
-import gradio as gr
-from langchain_community.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 transformers import TrOCRProcessor, TrOCRForConditionalGeneration
 from pdf2image import convert_from_path
+from PIL import Image
+import torch
 import os
 
-class LayoutLMv3OCR:
+# Klasse für die Textextraktion mit trOCR
+class TrOCR:
     def __init__(self):
-        # Lade den LayoutLMv3-Prozessor und das Modell für Token-Klassifikation
-        self.processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
-        self.model = AutoModelForTokenClassification.from_pretrained("microsoft/layoutlmv3-base")
+        # Laden des TrOCR Prozessors und Modells
+        self.processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
+        self.model = TrOCRForConditionalGeneration.from_pretrained("microsoft/trocr-base-stage1")
 
     def extract_text(self, pdf_path):
+        # Konvertiere PDF-Seiten in Bilder
         images = convert_from_path(pdf_path)
         text_pages = []
+        
         for image in images:
-            # Verarbeite die Bilddaten mit LayoutLMv3
-            inputs = self.processor(images=image, return_tensors="pt")
-            # Führe Vorhersagen durch
-            outputs = self.model(**inputs)
-            # Extrahiere den Text aus den Vorhersagen (falls dies vorgesehen ist)
-            text = self.processor.batch_decode(outputs.logits, skip_special_tokens=True)[0]
+            # Prozessiere das Bild mit trOCR
+            pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
+            generated_ids = self.model.generate(pixel_values)
+            text = self.processor.decode(generated_ids[0], skip_special_tokens=True)
             text_pages.append(text)
+        
         return text_pages
 
-ocr_tool = LayoutLMv3OCR()
+# OCR-Tool initialisieren
+ocr_tool = TrOCR()
 
+# Beispiel: Funktion zum Verarbeiten von PDFs und Beantworten von Fragen
 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)
-
-    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
+    # Extrahiere Text aus dem PDF
+    extracted_text = ocr_tool.extract_text(pdf_path)
+    # In diesem Beispiel geben wir den gesamten extrahierten Text zurück
+    # Hier kannst du Logik hinzufügen, um gezielt Fragen zu beantworten
+    return "\n".join(extracted_text)
 
+# Gradio Funktion zum Einlesen der Datei und Beantworten von Fragen
 def chatbot_response(pdf, question):
-    # Speichern der hochgeladenen Datei auf dem lokalen Dateisystem
     pdf_path = "uploaded_pdf.pdf"
-    
-    # Schreibe die PDF-Datei in eine lokale Datei
+    # Speichern der hochgeladenen PDF-Datei
     with open(pdf_path, "wb") as f:
         f.write(pdf.read())
     
-    extracted_text = ocr_tool.extr
+    # Extrahiere Text und beantworte die Frage
+    answer = process_pdf_and_query(pdf_path, question)
+    
+    # Nach der Verarbeitung die Datei löschen
+    os.remove(pdf_path)
+    
+    return answer
+
+# Beispiel für Gradio UI
+import gradio as gr
+
+with gr.Blocks() as demo:
+    pdf_input = gr.File(label="PDF-Datei hochladen")
+    question_input = gr.Textbox(label="Frage eingeben")
+    output = gr.Textbox(label="Antwort")
+    
+    pdf_input.upload(chatbot_response, inputs=[pdf_input, question_input], outputs=output)
+
+demo.launch()