-
Notifications
You must be signed in to change notification settings - Fork 0
/
nlpapp.py
39 lines (32 loc) · 1.25 KB
/
nlpapp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
#Aug/2021
import numpy as np
import pandas as pd
import streamlit as st
from streamlit import caching
from transformers import BartForConditionalGeneration, BartTokenizer, BartConfig
#App works off large BART mdel. Model is downloaded once the app runs and is cached.
#Loading the model and tokenizer for bart-large-cnn
@st.cache(allow_output_mutation=True)
def get_model():
tokenizer=BartTokenizer.from_pretrained('facebook/bart-large-cnn')
model=BartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn')
return tokenizer,model
tokenizer,model= get_model()
#Use pretrained model and tokenizer to produce summary tokens and then decode these summarized tokens
@st.cache(allow_output_mutation=True)
def summarizer(original_text):
inputs = tokenizer.batch_encode_plus([original_text],return_tensors='pt')
summary_ids = model.generate(inputs['input_ids'], early_stopping=True)
bart_summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return bart_summary
def main():
""" NLP Based App with Streamlit """
# Title
st.title("Summarize text")
# Summarization
message = st.text_area("Enter Text")
if st.button("Summarize"):
summary_result = summarizer(message)
st.success(summary_result)
if __name__ == '__main__':
main()