import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# preprocess input
# return input_ids matrix
tokenizer = AutoTokenizer.from_pretrained("SophieTr/fine-tune-Pegasus")
model = AutoModelForSeq2SeqLM.from_pretrained("SophieTr/fine-tune-Pegasus")

def preprocess(inp):
    input_ids = tokenizer(inp, return_tensors="pt").input_ids
    return input_ids
def predict(input_ids):
    outputs = model.generate(input_ids=input_ids)
    res = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
    return res

@app.route('/', methods=['POST', 'GET'])
def index():
    if request.method == 'POST':
        inp = request.form['content']
        inp_ids = preprocess(inp)
        summary = predict(inp_ids)
        return render_template('index.html', summary=summary)
    else:
        print("GETTING get")
        return render_template('index.html', summary="Nothing to summarize")


if __name__ == '__main__':
    st.title("Text summary with fine-tuned Pegasus model")
    with st.container():
        txt = st.text_area('Text to analyze', ' ')
        inp_ids = preprocess(txt)
        st.write('Summary:', predict(inp_ids))