Update app.py

app.py

@@ -10,12 +10,13 @@ def load_model(cp_path):
 @st.cache_resource
 def load_tokenizer(path):
     tokenizer = AutoTokenizer.from_pretrained(path)
+    return tokenizer
 
 cp_aug = 'minnehwg/finetune-newwiki-summarization-ver-augmented2'
 cp_org = 'minnehwg/finetune-newwiki-summarization-ver2'
 
 model_org = load_model(cp_org)
-model_aug  = AutoModelForSeq2SeqLM.from_pretrained(cp_aug)
+model_aug  = load_model(cp_aug)
 tokenizer = load_tokenizer("VietAI/vit5-base")
 
 def summarize(text, model, tokenizer, num_beams=4, device='cpu'):