Summarizer / app.py
ahabb's picture
Create app.py
395e99f verified
raw
history blame
969 Bytes
import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline, Seq2SeqTrainer, Seq2SeqTrainingArguments
model_path = 'T5_samsum'
# Load the model
model = AutoModelForSequenceClassification.from_pretrained(model_path)
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)
# Create the summarization pipeline
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
# Define the summarization function
def summarize_dialogue(dialogue):
summary = summarizer(dialogue, max_length=150, min_length=50, do_sample=False)
return summary[0]['summary_text']
# Create the Gradio interface
iface = gr.Interface(
fn=summarize_dialogue,
inputs=gr.Textbox(lines=10, placeholder="Enter the dialogue here..."),
outputs="text",
title="Dialogue Summarizer",
description="Enter a dialogue and this app will generate a summary using a pre-trained model."
)
# Launch the app
iface.launch()