import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
from PyPDF2 import PdfReader

# Models and tokenizers setup
models = {
    "Text Generator (Bloom)": {
        "model": AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m"),
        "tokenizer": AutoTokenizer.from_pretrained("bigscience/bloom-560m"),
    },
    "PDF Summarizer (T5)": {
        "model": AutoModelForSeq2SeqLM.from_pretrained("t5-small"),
        "tokenizer": AutoTokenizer.from_pretrained("t5-small", use_fast=False),  # Use the slow tokenizer
    },
    "Broken Answer (T0pp)": {
        "model": AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp"),
        "tokenizer": AutoTokenizer.from_pretrained("bigscience/T0pp", use_fast=False),  # Use the slow tokenizer
    },
}


# Function for text generation
def generate_text(model_choice, input_text, max_tokens, temperature, top_p):
    model_info = models[model_choice]
    tokenizer = model_info["tokenizer"]
    model = model_info["model"]

    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    outputs = model.generate(
        **inputs, max_length=max_tokens, num_beams=5, early_stopping=True, temperature=temperature, top_p=top_p
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Function for PDF summarization
def summarize_pdf(pdf_file, max_tokens, temperature, top_p):
    reader = PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    
    model_info = models["PDF Summarizer (T5)"]
    tokenizer = model_info["tokenizer"]
    model = model_info["model"]
    
    inputs = tokenizer("summarize: " + text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    outputs = model.generate(
        **inputs, max_length=max_tokens, num_beams=5, early_stopping=True, temperature=temperature, top_p=top_p
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Build Gradio interface
def launch_custom_app():
    with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
        gr.Markdown("<h1 style='text-align: center;'>💡 Multi-Model Assistant</h1>")
        gr.Markdown("<p style='text-align: center;'>Switch between text generation, PDF summarization, or quirky broken answers!</p>")
        
        with gr.Tabs():
            # Tab for Text Generation
            with gr.Tab("Text Generator"):
                model_choice = gr.Dropdown(choices=list(models.keys()), label="Choose a Model", value="Text Generator (Bloom)")
                input_text = gr.Textbox(label="Enter Text")
                max_tokens = gr.Slider(minimum=10, maximum=512, value=150, step=10, label="Max Tokens")
                temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
                output_text = gr.Textbox(label="Generated Text", interactive=False)
                generate_button = gr.Button("Generate Text")
                
                generate_button.click(
                    generate_text, 
                    inputs=[model_choice, input_text, max_tokens, temperature, top_p],
                    outputs=output_text
                )
            
            # Tab for PDF Summarization
            with gr.Tab("PDF Summarizer"):
                pdf_file = gr.File(label="Upload a PDF File", file_types=[".pdf"])
                max_tokens_pdf = gr.Slider(minimum=10, maximum=512, value=150, step=10, label="Max Tokens")
                temperature_pdf = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
                top_p_pdf = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
                summary_output = gr.Textbox(label="PDF Summary", interactive=False)
                summarize_button = gr.Button("Summarize PDF")
                
                summarize_button.click(
                    summarize_pdf, 
                    inputs=[pdf_file, max_tokens_pdf, temperature_pdf, top_p_pdf],
                    outputs=summary_output
                )
            
            # Tab for Broken Model
            with gr.Tab("Broken Answers"):
                broken_input = gr.Textbox(label="Enter Text")
                broken_max_tokens = gr.Slider(minimum=10, maximum=512, value=150, step=10, label="Max Tokens")
                broken_temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
                broken_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
                broken_output = gr.Textbox(label="Broken Model Output", interactive=False)
                broken_button = gr.Button("Generate Broken Answer")
                
                broken_button.click(
                    lambda text, max_tokens, temp, top_p: generate_text("Broken Answer (T0pp)", text, max_tokens, temp, top_p),
                    inputs=[broken_input, broken_max_tokens, broken_temperature, broken_top_p],
                    outputs=broken_output
                )
        
        demo.launch()

# Launch the app
launch_custom_app()