Sanj004 commited on
Commit
a5cb665
·
verified ·
1 Parent(s): 03f8910

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +34 -0
app.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
+ from datasets import load_dataset
4
+ import torch
5
+
6
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
7
+
8
+ SAVED_MODEL_PATH = './bart_base_full_finetune_save'
9
+ model_name = "facebook/bart-base"
10
+ model = AutoModelForSeq2SeqLM.from_pretrained(SAVED_MODEL_PATH).to(device)
11
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
12
+
13
+ dataset = load_dataset("samsum")
14
+
15
+ train_data = dataset["train"]
16
+ validation_data = dataset["validation"]
17
+ test_data = dataset["test"]
18
+
19
+ def summarize(tokenizer, model, text):
20
+ inputs = tokenizer(f"Summarize dialogue >>\n {text}", return_tensors="pt", max_length=1000, truncation=True, padding="max_length").to(device)
21
+ summary_ids = model.generate(inputs.input_ids, num_beams=4, max_length=100, early_stopping=True)
22
+ summary = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids]
23
+ return summary[0]
24
+
25
+ def summarize_dialogue(text):
26
+ return summarize(tokenizer, model, text)
27
+
28
+ iface = gr.Interface(
29
+ fn=summarize_dialogue,
30
+ inputs=gr.inputs.Textbox(lines=10, label="Input Dialogue"),
31
+ outputs=gr.outputs.Textbox(label="Generated Summary")
32
+ )
33
+
34
+ iface.launch()