22f1001555
Deploy conversation summarizer space
065beff
Raw
History Blame Contribute Delete
1.95 kB
from transformers import T5ForConditionalGeneration, T5Tokenizer
import gradio as gr
model = T5ForConditionalGeneration.from_pretrained("abhi-codes/finetuned_flank_t5_for_summarization")
tokenizer = T5Tokenizer.from_pretrained("abhi-codes/finetuned_flank_t5_for_summarization")
examples = [
["""Tom: Did you submit the report?
Anika: Not yet, I'm fixing the charts.
Tom: The deadline is 5 pm.
Anika: I know. I'll send it by 4:30.
Tom: Great, please copy me on the email."""],
["""Nora: Are you picking up the groceries today?
Eli: Yes, after work.
Nora: Please get milk, eggs, and bread.
Eli: Got it. Anything else?
Nora: Bananas if they look fresh.
Eli: Okay, I'll be home around 6:30"""],
["""Priya: Did you call the dentist?
Karan: Yes, they had an opening tomorrow at 11.
Priya: Great. Did you book it?
Karan: Yes, I confirmed it.
Priya: Thanks. I'll leave work early to go."""]
]
def summarize(input):
input = "summarize: "+ input
model_inputs = tokenizer(input, return_tensors="pt", max_length=512, truncation=True,padding = 'max_length')
summary_ids = model.generate(
input_ids=model_inputs["input_ids"],
attention_mask=model_inputs["attention_mask"],
max_new_tokens=128,
num_beams=4,
no_repeat_ngram_size=3
)
return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
demo = gr.Interface(
fn=summarize,
inputs=[
gr.Textbox(
lines=8,
label="Dialogue",
placeholder="Paste a conversation here"
)],
outputs=[
gr.Textbox(
lines=2,
label="Summary"
),
],
title="Dialogue Summarizer",
description=(
"Enter a chat-style conversation and the model will generate a short summary. "
"For best results, write each message on a new line with the speaker name."
),
examples=examples,
flagging_mode="never"
)
demo.launch()