from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
import gradio as gr
from PyPDF2 import PdfReader

# Modell und Tokenizer für Fragebeantwortung laden
model_name = "deepset/roberta-base-squad2"
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)

# Funktion zum Extrahieren von Text aus der PDF
def extract_text_from_pdf(pdf_path):
    reader = PdfReader(pdf_path)
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    return text

# Funktion für die Fragebeantwortung
def chatbot_response(pdf_path, question):
    # PDF-Text extrahieren
    context = extract_text_from_pdf(pdf_path)
    
    # Frage beantworten
    try:
        result = qa_pipeline(question=question, context=context)
        return result['answer']
    except Exception as e:
        return f"Fehler bei der Beantwortung: {e}"

# Gradio-Interface erstellen
pdf_input = gr.File(label="PDF-Datei hochladen", type="filepath")
question_input = gr.Textbox(label="Frage eingeben", placeholder="Stelle eine Frage zu dem PDF-Dokument")
response_output = gr.Textbox(label="Antwort")

# Gradio-Interface
interface = gr.Interface(
    fn=chatbot_response,
    inputs=[pdf_input, question_input],
    outputs=response_output,
    title="PDF-Fragebeantwortung auf Deutsch",
    description="Lade eine PDF-Datei hoch und stelle Fragen zu ihrem Inhalt. Antworten basieren nur auf den PDF-Inhalten."
)

if __name__ == "__main__":
    interface.launch()