import torch
import gradio as gr
from transformers import pipeline

# Check CUDA availability
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU Device:", torch.cuda.get_device_name(0))


# Initialize the summarization pipeline with GPU support
device = 0 if torch.cuda.is_available() else -1
text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16, device=device)

# Function to summarize text
def summary(input):
    output = text_summary(input, max_length=130, min_length=30, do_sample=False)  # Fixed pipeline name
    return output[0]['summary_text']


# Create the Gradio Interface
gr.close_all()
demo = gr.Interface(
    fn=summary,
    inputs=[gr.Textbox(label="Input Text to Summarize", lines=10)],
    outputs=[gr.Textbox(label="Summarized Text", lines=6)],
    title="A.C. Text Summarizer",
    description="This application will be used to create summarized text"
)

# Launch the interface
demo.launch(share=True)