import gradio as gr
from transformers import pipeline
import torch

# Check if GPU is available for FP16 inference
device = 0 if torch.cuda.is_available() else -1
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# Load Instruction-Tuned Summarization Model
summarization = pipeline(
    "text2text-generation",
    model="google/flan-t5-base",
    device=device,
    torch_dtype=torch_dtype
)

def summarize_text(text):
    # Ensure the input is not too short for summarization
    if len(text.split()) < 30:
        return "Please provide a longer text (at least 30 words) for better summarization."

    # Prompt-based instruction
    prompt = f"Summarize the following text in 3 to 5 clear and concise sentences, highlighting the main points and key takeaways:\n\n{text}"
    
    output = summarization(prompt, max_length=256, do_sample=False)
    return output[0]['generated_text']

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Instruction-based Text Summarization (FLAN-T5)")

    with gr.Tab("Text Summarization"):
        summary_input = gr.Textbox(label="Text to Summarize", lines=10, placeholder="Paste long text here...")
        summary_btn = gr.Button("Summarize")
        summary_output = gr.Textbox(label="Summary")
        summary_btn.click(fn=summarize_text, inputs=summary_input, outputs=summary_output)

# Launch App
demo.launch()