SophieTr commited on
Commit
a214965
1 Parent(s): 1e7e7cc

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from flask import Flask, request, render_template, redirect
3
+ from transformers import BartTokenizer, BartForConditionalGeneration
4
+
5
+ # preprocess input
6
+ # return input_ids matrix
7
+ tokenizer = BartTokenizer.from_pretrained("sshleifer/distilbart-cnn-6-6")
8
+ model = BartForConditionalGeneration.from_pretrained("sshleifer/distilbart-cnn-6-6")
9
+
10
+ def preprocess(inp):
11
+ input_ids = tokenizer(inp, return_tensors="pt").input_ids
12
+ return input_ids
13
+ def predict(input_ids):
14
+ outputs = model.generate(input_ids=input_ids)
15
+ res = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
16
+ return res
17
+
18
+ @app.route('/', methods=['POST', 'GET'])
19
+ def index():
20
+ if request.method == 'POST':
21
+ inp = request.form['content']
22
+ inp_ids = preprocess(inp)
23
+ summary = predict(inp_ids)
24
+ return render_template('index.html', summary=summary)
25
+ else:
26
+ print("GETTING get")
27
+ return render_template('index.html', summary="Nothing to summarize")
28
+
29
+
30
+ if __name__ == '__main__':
31
+ st.title("Text summary with fine-tuned Pegasus model")
32
+ with st.container():
33
+ txt = st.text_area('Text to analyze', ' ')
34
+ inp_ids = preprocess(txt)
35
+ st.write('Summary:', predict(inp_ids))
36
+