SophieTr/TextSummarizationDemo

1import streamlit as st2 3from transformers import AutoTokenizer, AutoModelForSeq2SeqLM4 5tokenizer = AutoTokenizer.from_pretrained("QuickRead/pegasus-reddit-7e05-new")6 7model = AutoModelForSeq2SeqLM.from_pretrained("QuickRead/pegasus-reddit-7e05-new")8 9def preprocess(inp):10    input_ids = tokenizer(inp, return_tensors="pt").input_ids11    return input_ids12def predict(input_ids):13    outputs = model.generate(input_ids=input_ids)14    res = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]15    return res16 17if __name__ == '__main__':18    st.title("Text summary with fine-tuned Pegasus model")19    with st.container():20        txt = st.text_area('Text to analyze', ' ')21        inp_ids = preprocess(txt)22        st.write('Summary:', predict(inp_ids))23