Update app.py

app.py

@@ -4,19 +4,18 @@ import nltk
 import math
 import torch
 
-model_name = "fabiochiu/t5-base-medium-title-generation"
+model_name = "Nopphakorn/mt5-small-thaisum-512-title"
 max_input_length = 512
 
-st.header("Generate candidate titles for articles")
+st.header("Generate headline titles for Thai news")
 
-st_model_load = st.text('Loading title generator model...')
+st_model_load = st.text('Loading headlines summarizer model...')
 
 @st.cache(allow_output_mutation=True)
 def load_model():
     print("Loading model...")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    nltk.download('punkt')
     print("Model loaded!")
     return tokenizer, model
 
@@ -81,14 +80,14 @@ def generate_title():
     }
 
     # compute predictions
-    outputs = model.generate(**inputs, do_sample=True, temperature=temperature)
+    outputs = model.generate(**inputs, do_sample=True, temperature=temperature, max_new_tokens=64)
     decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)
     predicted_titles = [nltk.sent_tokenize(decoded_output.strip())[0] for decoded_output in decoded_outputs]
 
     st.session_state.titles = predicted_titles
 
 # generate title button
-st_generate_button = st.button('Generate title', on_click=generate_title)
+st_generate_button = st.button('Generate headlines', on_click=generate_title)
 
 # title generation labels
 if 'titles' not in st.session_state:
@@ -96,6 +95,6 @@ if 'titles' not in st.session_state:
 
 if len(st.session_state.titles) > 0:
     with st.container():
-        st.subheader("Generated titles")
+        st.subheader("Generated headlines")
         for title in st.session_state.titles:
             st.markdown("__" + title + "__")