Javedalam's picture
Update app.py
1ddbd5c verified
raw
history blame contribute delete
No virus
1.26 kB
import gradio as gr
from transformers import pipeline
# Step 3: Define the summarization function for multiple models
summarizers = {
"BART (facebook/bart-large-cnn)": pipeline("summarization", model="facebook/bart-large-cnn"),
"T5 (t5-small)": pipeline("summarization", model="t5-small"),
"Pegasus (google/pegasus-xsum)": pipeline("summarization", model="google/pegasus-xsum"),
"DistilBART (sshleifer/distilbart-cnn-12-6)": pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
}
def summarize(text, model_name):
summarizer = summarizers[model_name]
summary = summarizer(text, max_length=150, min_length=40, do_sample=False)
return summary[0]['summary_text']
# Step 4: Create the Gradio interface
description = """
Summarize text using various models from Hugging Face:
- BART (facebook/bart-large-cnn)
- T5 (t5-small)
- Pegasus (google/pegasus-xsum)
- DistilBART (sshleifer/distilbart-cnn-12-6)
"""
iface = gr.Interface(
fn=summarize,
inputs=[
gr.Textbox(lines=10, label="Input Text"),
gr.Dropdown(choices=list(summarizers.keys()), label="Choose Model")
],
outputs="textbox",
title="Text Summarizer",
description=description
)
# Step 5: Launch the interface
iface.launch()