import gradio as gr
import torch
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import spaces  # 1. IMPORT SPACES FOR ZEROGPU

transformers.logging.set_verbosity_error()

MODEL_NAME = "google/pegasus-pubmed"

# 2. FORCE CUDA: ZeroGPU virtualizes the GPU, so we can hardcode "cuda" safely
device = "cuda" 

print("Loading model into ZeroGPU memory...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)

# 3. ADD DECORATOR: This tells Hugging Face to assign a GPU just for this function
@spaces.GPU
def summarize_medical_abstract(abstract_text):
    if not abstract_text or len(abstract_text.strip()) < 50:
        return "Error: Input text is too short or invalid. Please provide a full abstract."

    inputs = tokenizer(
        abstract_text, 
        truncation=True, 
        padding="longest", 
        max_length=1024, 
        return_tensors="pt"
    ).to(device)
    
    with torch.no_grad():
        summary_ids = model.generate(
            inputs["input_ids"],
            max_length=150,
            min_length=40,
            num_beams=4,
            length_penalty=2.0,
            no_repeat_ngram_size=3,
            do_sample=False,
            early_stopping=True
        )
    
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)

# UI Configuration
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🩺 Medical Research Paper Abstract Summarizer")
    gr.Markdown("**Developed by:** Debarghya Bhowmick | **Guide:** Dr. Tohida Rehman (JU)\n*Powered by Pegasus-PubMed (ZeroGPU)*")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                lines=12, 
                placeholder="Paste the medical research abstract here...", 
                label="Input: Clinical Abstract"
            )
            submit_btn = gr.Button("Generate Summary", variant="primary")
            
        with gr.Column():
            output_text = gr.Textbox(
                lines=12, 
                label="Output: Extractive Summary",
                interactive=False
            )
            
    submit_btn.click(
        fn=summarize_medical_abstract,
        inputs=input_text,
        outputs=output_text
    )

demo.launch()