import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 1. Load the model and tokenizer directly into the app memory
model_name = "sshleifer/distilbart-cnn-12-6"
print("Loading model... this takes a moment when the space wakes up.")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 2. Define the logic
def ai_summarizer_interface(input_text):
    if not input_text.strip():
        return "Error: Please input text to analyze."
    
    inputs = tokenizer(input_text, max_length=1024, truncation=True, return_tensors="pt")
    
    with torch.no_grad():
        summary_ids = model.generate(
            inputs["input_ids"], 
            num_beams=4, 
            max_length=150, 
            min_length=40, 
            early_stopping=True
        )
        
    generated_summary = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)[0]
    return generated_summary

# 3. Construct the UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📑 Automated Text Summarization Engine")
    gr.Markdown("### Developed for Jadavpur University Research Internship by Debarghya Bhowmick (Under the guidance of respected Dr. Tohida Rehman)")
    
    with gr.Row():
        with gr.Column():
            input_box = gr.Textbox(
                lines=12, 
                label="Source Document / Research Paper Text", 
                placeholder="Paste long-form text here..."
            )
            submit_btn = gr.Button("Generate Abstractive Summary", variant="primary")
        with gr.Column():
            output_box = gr.Textbox(
                lines=6, 
                label="System Generated Summary (Abstractive Baseline)", 
                interactive=False
            )
            gr.Markdown("**Analysis Focus:** Use this sandbox prototype to cross-verify the output against the source text to log structural hallucinations or fact-distortion patterns.")

    submit_btn.click(fn=ai_summarizer_interface, inputs=input_box, outputs=output_box)

# 4. Launch the app (no share=True needed for HF Spaces)
demo.launch()