ai-lab-macree's picture
Added a slider for input length
190132f
raw
history blame
945 Bytes
'''
A script that uses gradio to make a web app that takes in a dialog and outputs a summary.
Import a huggingface model.
'''
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
MODEL_NAME = "Firefly777a/bart-large-samsum-candice-set-summary-3ep-3"
# Import the model
summarizer = pipeline("summarization", model=MODEL_NAME, tokenizer=MODEL_NAME)
# tokenizer = AutoTokenizer.from_pretrained(PRETRAIN_MODEL_NAME)
# Define the function that will be used by the web app
def summarize_dialog(dialog, max_length):
# Summarize the dialog
summary = summarizer(dialog, max_length=max_length, min_length=10, do_sample=False)
# Return the summary
return summary[0]['summary_text']
# Define the web app & launch
gr.Interface(fn=summarize_dialog,
inputs=["text",
gr.Slider(32, 256, value=128)
],
outputs="text").launch()