charlie572 commited on
Commit
a258be0
·
1 Parent(s): bc2836f

Get summarizer to work

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -1,10 +1,15 @@
1
  import gradio as gr
2
- from transformers import pipeline
3
 
4
  def summarize(text):
5
- summarizer = pipeline("summarization")
6
- result = summarizer(text)
7
- return result["summary_text"]
 
 
 
 
 
8
 
9
  iface = gr.Interface(fn=summarize, inputs="text", outputs="text")
10
  iface.launch()
 
1
  import gradio as gr
2
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
 
4
  def summarize(text):
5
+ checkpoint = "sshleifer/distilbart-cnn-12-6"
6
+ tokenizer = AutoTokenizer.from_pretrained(checkpoint)
7
+ model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
8
+
9
+ inputs = tokenizer(text, truncation=True, return_tensors="pt").input_ids
10
+ outputs = model.generate(inputs, max_new_tokens=100, do_sample=False)
11
+
12
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
13
 
14
  iface = gr.Interface(fn=summarize, inputs="text", outputs="text")
15
  iface.launch()