import streamlit as st

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("QuickRead/pegasus-reddit-7e05-new")

model = AutoModelForSeq2SeqLM.from_pretrained("QuickRead/pegasus-reddit-7e05-new")

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

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))