jknebel commited on
Commit
812e1bd
1 Parent(s): 1e6cafd

adapt for better working

Browse files
Files changed (2) hide show
  1. app.py +24 -16
  2. requirements.txt +2 -1
app.py CHANGED
@@ -1,25 +1,33 @@
1
  from transformers import pipeline
2
  import gradio as gr
3
 
4
- model_name_converter = {"bart_large":"juliosocher/bart-large-cnn-finetuned-scientific-articles",
5
- "mt5-small-finetuned-mt5":"jacks392/mt5-small-finetuned-mt5",
6
- "facebook": "facebook/bart-large-cnn",
7
- "google" : "google/pegasus-xsum"
8
- }
9
- def predict(prompt,model_name, max_length):
10
- if model_name ==None:
11
- model_name = "google/pegasus-xsum"
 
 
 
 
 
 
 
12
  else:
13
- model_name = model_name_converter[model_name]
14
- model = pipeline("summarization",model = model_name)
15
- summary = model(prompt,max_length)[0]["summary_text"]
16
  return summary
17
-
18
- options_1 = list(model_name_converter.keys())
 
19
  with gr.Blocks() as demo:
20
  drop_down = gr.Dropdown(choices=options_1, label="model")
21
- textbox = gr.Textbox(placeholder = "Enter text block to summarize", lines = 4)
22
- length=gr.Number(value = 200, label="the max number of characher for summerized")
23
- gr.Interface(fn=predict, inputs=[textbox, drop_down, length], outputs = "text")
24
 
25
  demo.launch()
 
1
  from transformers import pipeline
2
  import gradio as gr
3
 
4
+
5
+ MODELS = {
6
+ "gsarti": pipeline("summarization", model="gsarti/it5-base-wiki-summarization"),
7
+ "facebook": pipeline("summarization", model="facebook/bart-large-cnn"),
8
+ "lincoln": pipeline(
9
+ "summarization", model="lincoln/mbart-mlsum-automatic-summarization"
10
+ ),
11
+ "google": pipeline("summarization", model="google/pegasus-large"),
12
+ "t5-small": pipeline("summarization", model="t5-small"),
13
+ }
14
+
15
+
16
+ def predict(prompt, model_name, max_length):
17
+ if model_name is None:
18
+ model = MODELS["t5-small"]
19
  else:
20
+ model = MODELS[model_name]
21
+ prompt = prompt.replace("\n", " ")
22
+ summary = model(prompt, max_length)[0]["summary_text"]
23
  return summary
24
+
25
+
26
+ options_1 = list(MODELS.keys())
27
  with gr.Blocks() as demo:
28
  drop_down = gr.Dropdown(choices=options_1, label="model")
29
+ textbox = gr.Textbox(placeholder="Enter text block to summarize", lines=4)
30
+ length = gr.Number(value=100, label="the max number of characher for summerized")
31
+ gr.Interface(fn=predict, inputs=[textbox, drop_down, length], outputs="text")
32
 
33
  demo.launch()
requirements.txt CHANGED
@@ -1,4 +1,5 @@
1
  gradio
2
  transformers
3
  tensorflow
4
- tf-keras
 
 
1
  gradio
2
  transformers
3
  tensorflow
4
+ tf-keras
5
+ sentencepiece