Summarizer / Text_Summary_app.py
PRSHNTKUMR's picture
Update Text_Summary_app.py
be43a82 verified
import os
import io
from IPython.display import Image, display, HTML
from PIL import Image
import base64
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")
import json
def summarize_text(input_text):
"""Summarizes the given input text.
Args:
input_text (str): The text to be summarized.
Returns:
dict: A dictionary containing the summary under the 'summary' key.
"""
# Tokenize and truncate input if necessary
max_length = tokenizer.model_max_length
inputs = tokenizer(input_text, truncation=True, max_length=max_length, return_tensors="pt")
# Generate summary
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 as a dictionary
return {"summary": summary_text}
def generate_summary(input):
output = summarize_text(input)
return output
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(share=True, server_port=int(os.environ['PORT']))