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
import os

class LayoutLMv3OCR:
    def __init__(self):
        self.processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
        # Ändere AutoModelForSeq2SeqLM zu AutoModelForTokenClassification
        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:
            # Bilder werden für die OCR-Prozesse vorbereitet
            inputs = self.processor(images=image, return_tensors="pt")
            # Modell wird zur Textextraktion genutzt
            outputs = self.model(**inputs)
            # Hier wird der dekodierte Text extrahiert
            text = self.processor.batch_decode(outputs.logits, skip_special_tokens=True)[0]
            text_pages.append(text)
        return text_pages

ocr_tool = LayoutLMv3OCR()

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

def chatbot_response(pdf, question):
    pdf_path = "uploaded_pdf.pdf"
    pdf.save(pdf_path)
    extracted_text = ocr_tool.extract_text(pdf_path)
    answer = process_pdf_and_query(pdf_path, question)
    os.remove(pdf_path)
    return answer

pdf_input = gr.inputs.File(label="PDF-Datei hochladen")
question_input = gr.inputs.Textbox(label="Frage eingeben")
response_output = gr.outputs.Textbox(label="Antwort")

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.launch()