Update app.py

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import nltk
-import math
 import torch
 
 model_name = "afnanmmir/t5-base-abstract-to-plain-language-1"
@@ -13,15 +12,14 @@ st.header("Generate summaries")
 
 st_model_load = st.text('Loading summary generator model...')
 
-# @st.cache(allow_output_mutation=True)
-@st.cache_data
-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
+# # @st.cache(allow_output_mutation=True)
+# @st.cache_data
+# def load_model():
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+nltk.download('punkt')
+print("Model loaded!")
 
 tokenizer, model = load_model()
 st.success('Model loaded!')