Update app.py

app.py

@@ -1,15 +1,18 @@
 import transformers
 import streamlit as st
 
-from transformers import AutoTokenizer, AutoModelWithLMHead
-  
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+tokenizer = AutoTokenizer.from_pretrained("UNIST-Eunchan/bart-dnc-booksum")
+
 tokenizer = AutoTokenizer.from_pretrained("gpt2-large")
 @st.cache
 def load_model(model_name):
-    model = AutoModelWithLMHead.from_pretrained("gpt2-large")
+
+    model = AutoModelForSeq2SeqLM.from_pretrained("UNIST-Eunchan/bart-dnc-booksum")
     return model
 
-model = load_model("gpt2-large")
+model = load_model("UNIST-Eunchan/bart-dnc-booksum")
 
 def infer(input_ids, max_length, temperature, top_k, top_p):
 
@@ -20,9 +23,11 @@ def infer(input_ids, max_length, temperature, top_k, top_p):
         top_k=top_k,
         top_p=top_p,
         do_sample=True,
-        num_return_sequences=1
+        num_return_sequences=1,
+        num_beams=4,
+        no_repeat_ngram_size=2
     )
-
+    
     return output_sequences
 default_value = "See how a modern neural network auto-completes your text 🤗 This site, built by the Hugging Face team, lets you write a whole document directly from your browser, and you can trigger the Transformer anywhere using the Tab key. Its like having a smart machine that completes your thoughts 😀 Get started by typing a custom snippet, check out the repository, or try one of the examples. Have fun!"
 
@@ -30,11 +35,11 @@ default_value = "See how a modern neural network auto-completes your text 🤗 T
 st.title("Write with Transformers 🦄")
 st.write("The almighty king of text generation, GPT-2 comes in four available sizes, only three of which have been publicly made available. Feared for its fake news generation capabilities, it currently stands as the most syntactically coherent model. A direct successor to the original GPT, it reinforces the already established pre-training/fine-tuning killer duo. From the paper: Language Models are Unsupervised Multitask Learners by Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei and Ilya Sutskever.")
 
-sent = st.text_area("Text", default_value, height = 275)
-max_length = st.sidebar.slider("Max Length", min_value = 10, max_value=30)
+sent = st.text_area("Text", default_value, height = 550)
+max_length = st.sidebar.slider("Max Length", min_value = 10, max_value=256)
 temperature = st.sidebar.slider("Temperature", value = 1.0, min_value = 0.0, max_value=1.0, step=0.05)
 top_k = st.sidebar.slider("Top-k", min_value = 0, max_value=5, value = 0)
-top_p = st.sidebar.slider("Top-p", min_value = 0.0, max_value=1.0, step = 0.05, value = 0.9)
+top_p = st.sidebar.slider("Top-p", min_value = 0.0, max_value=1.0, step = 0.05, value = 0.92)
 
 encoded_prompt = tokenizer.encode(sent, add_special_tokens=False, return_tensors="pt")
 if encoded_prompt.size()[-1] == 0: