import gradio as gr
from transformers import pipeline
import textwrap

# Load the question-answering pipeline
qa_pipeline = pipeline("question-answering")

def chatbot(document, question):
    # Define a set of greeting phrases
    greetings = ["hi", "hello", "hey", "greetings", "what's up", "howdy"]

    # Check if the input question is a greeting
    question_lower = question.lower().strip()
    if question_lower in greetings or any(question_lower.startswith(greeting) for greeting in greetings):
        return "Hello! How can I assist you with the document today?"
    
    # Otherwise, handle the question using the QA pipeline
    result = qa_pipeline(question=question, context=document)
    
    # Wrap the answer to ensure it is 3 to 4 lines long
    wrapped_answer = textwrap.fill(result['answer'], width=70)
    
    # Split the wrapped answer into lines and limit it to 3 to 4 lines
    answer_lines = wrapped_answer.split('\n')
    limited_answer = '\n'.join(answer_lines[:4])
    
    return limited_answer

interface = gr.Interface(
    fn=chatbot,
    inputs=[
        gr.components.Textbox(lines=20, placeholder="Paste your document here..."),
        gr.components.Textbox(lines=2, placeholder="Ask a question about the document or say hello...")
    ],
    outputs="text",
    title="Document Chatbot",
    description="Upload a document and ask questions about its content or just say hello."
)

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