Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline, Seq2SeqTrainer, Seq2SeqTrainingArguments | |
model_path = 'T5_samsum' | |
# Load the model | |
model = AutoModelForSeq2SeqLM.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() | |