Update app.py

app.py

@@ -1,64 +1,65 @@
-from transformers import TrOCRProcessor, TrOCRForConditionalGeneration
+import gradio as gr
+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 PIL import Image
-import torch
 import os
 
-# Klasse für die Textextraktion mit trOCR
-class TrOCR:
+class LayoutLMv3OCR:
     def __init__(self):
-        # Laden des TrOCR Prozessors und Modells
-        self.processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
-        self.model = TrOCRForConditionalGeneration.from_pretrained("microsoft/trocr-base-stage1")
+        self.processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
+        self.model = AutoModelForTokenClassification.from_pretrained("microsoft/layoutlmv3-base")
 
     def extract_text(self, pdf_path):
-        # Konvertiere PDF-Seiten in Bilder
         images = convert_from_path(pdf_path)
         text_pages = []
-        
         for image in images:
-            # 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)
+            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 initialisieren
-ocr_tool = TrOCR()
+ocr_tool = LayoutLMv3OCR()
 
-# Beispiel: Funktion zum Verarbeiten von PDFs und Beantworten von Fragen
 def process_pdf_and_query(pdf_path, question):
-    # 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)
+    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
 
-# Gradio Funktion zum Einlesen der Datei und Beantworten von Fragen
 def chatbot_response(pdf, question):
     pdf_path = "uploaded_pdf.pdf"
-    # Speichern der hochgeladenen PDF-Datei
-    with open(pdf_path, "wb") as f:
-        f.write(pdf.read())
-    
-    # Extrahiere Text und beantworte die Frage
+    pdf.save(pdf_path)
+    extracted_text = ocr_tool.extract_text(pdf_path)
     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
+# Ändere 'inputs' und 'outputs' zur neuen Gradio API
+pdf_input = gr.File(label="PDF-Datei hochladen")
+question_input = gr.Textbox(label="Frage eingeben")
+response_output = gr.Textbox(label="Antwort")
 
-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)
+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."
+)
 
-demo.launch()
+if __name__ == "__main__":
+    interface.launch()