Summarization / app.py
USF00's picture
Upload 2 files
f6ca911 verified
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()