| 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() |
|
|