AI_Apps / app.py
PRSHNTKUMR's picture
Create app.py
1fea30f verified
import gradio as gr
import os
from transformers import pipeline, AutoTokenizer
# Initialize summarizer and tokenizer
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", tokenizer="sshleifer/distilbart-cnn-12-6")
tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
def summarize_text(input_text):
"""Summarizes the given input text."""
max_length = tokenizer.model_max_length
inputs = tokenizer(input_text, truncation=True, max_length=max_length, return_tensors="pt")
summary_ids = summarizer.model.generate(inputs.input_ids, max_length=50, min_length=10, do_sample=False)
summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return {"summary": summary_text}
def generate_summary(input):
output = summarize_text(input)
return output["summary"] # Return the summary directly
gr.close_all()
demo = gr.Interface(
fn=generate_summary,
inputs=[gr.Textbox(label="Text to summarize", lines=6)],
outputs=[gr.Textbox(label="Summary", lines=3)],
title="Text Summarization",
description="Summarize text using the 'shleifer/distilbart-cnn-12-6' language model.",
)
demo.launch()