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

MODEL_NAME = "facebook/bart-large-cnn"

# load model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

def summarize(text):
    if not text:
        return "Please enter text"

    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024).to(device)

    summary_ids = model.generate(
        inputs["input_ids"],
        max_length=150,
        min_length=50,
        num_beams=4
    )

    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary

app = gr.Interface(
    fn=summarize,
    inputs=gr.Textbox(lines=10, placeholder="Enter long text here..."),
    outputs="text",
    title="LITVISION Summarizer",
    description="AI-powered book and text summarization"
)

app.launch()