Update app.py

app.py

@@ -1,84 +1,64 @@
 import gradio as gr
-from langchain.chains import RetrievalQA
 from langchain.vectorstores import Chroma
-from langchain.document_loaders import PyPDFLoader
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.schema import Document
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-
-# OCR-Ersatz: LayoutLMv3 für Textextraktion aus PDFs
-from transformers import LayoutLMv3Processor
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from transformers import LayoutLMv3Processor, AutoModelForSeq2SeqLM
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
 from pdf2image import convert_from_path
-from PIL import Image
-import torch
+import os
 
 class LayoutLMv3OCR:
     def __init__(self):
         self.processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
-        self.model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/layoutlmv3-base-finetuned", num_labels=2)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/layoutlmv3-base")
 
     def extract_text(self, pdf_path):
-        pages = convert_from_path(pdf_path)
-        extracted_texts = []
-        for page in pages:
-            encoding = self.processor(images=page, return_tensors="pt")
-            outputs = self.model(**encoding)
-            logits = outputs.logits
-            predictions = torch.argmax(logits, dim=-1).squeeze()
-            tokens = self.processor.tokenizer.convert_ids_to_tokens(encoding.input_ids[0])
-            page_text = " ".join([token for token, pred in zip(tokens, predictions) if pred == 1])
-            extracted_texts.append(page_text)
-        return extracted_texts
+        images = convert_from_path(pdf_path)
+        text_pages = []
+        for image in images:
+            inputs = self.processor(images=image, return_tensors="pt")
+            outputs = self.model.generate(**inputs)
+            text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
+            text_pages.append(text)
+        return text_pages
 
-# Initialisiere OCR
 ocr_tool = LayoutLMv3OCR()
 
-# Embeddings und LLM konfigurieren
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-model_name = "google/flan-t5-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+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)
 
-def flan_generate(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
-    outputs = model.generate(**inputs, max_length=512)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    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
 
-def process_pdf_and_create_rag(pdf_path):
+def chatbot_response(pdf, question):
+    pdf_path = "uploaded_pdf.pdf"
+    pdf.save(pdf_path)
     extracted_text = ocr_tool.extract_text(pdf_path)
-    documents = []
-    for page_num, text in enumerate(extracted_text, start=1):
-        doc = Document(page_content=text, metadata={"page": page_num})
-        documents.append(doc)
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    split_docs = text_splitter.split_documents(documents)
-    vector_store = Chroma.from_documents(split_docs, embedding=embeddings)
-    retriever = vector_store.as_retriever()
-    qa_chain = RetrievalQA(retriever=retriever, combine_documents_chain=flan_generate)
-    return qa_chain
+    answer = process_pdf_and_query(pdf_path, question)
+    os.remove(pdf_path)
+    return answer
 
-def chatbot_response(pdf_file, question):
-    qa_chain = process_pdf_and_create_rag(pdf_file.name)
-    response = qa_chain.run(question)
-    relevant_pages = set()
-    for doc in qa_chain.retriever.get_relevant_documents(question):
-        relevant_pages.add(doc.metadata.get("page", "Unbekannt"))
-    page_info = f" (Referenz: Seite(n) {', '.join(map(str, relevant_pages))})"
-    return response + page_info
+pdf_input = gr.inputs.File(label="PDF-Datei hochladen")
+question_input = gr.inputs.Textbox(label="Frage eingeben")
+response_output = gr.outputs.Textbox(label="Antwort")
 
-def gradio_interface():
-    pdf_input = gr.File(label="PDF hochladen")
-    question_input = gr.Textbox(label="Ihre Frage", placeholder="Geben Sie Ihre Frage hier ein...")
-    response_output = gr.Textbox(label="Antwort")
-    interface = gr.Interface(
-        fn=chatbot_response,
-        inputs=[pdf_input, question_input],
-        outputs=response_output,
-        title="RAG Chatbot (Deutsch)"
-    )
-    return interface
+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 = gradio_interface()
     interface.launch()