Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,6 +5,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import nltk
 import math
 import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 model_name = "InsafQ/title_creator"
 max_input_length = 512
@@ -16,8 +17,8 @@ st_model_load = st.text('Loading title generator model...')
 @st.cache_data()
 def load_model():
     print("Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+    tokenizer = T5Tokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = T5ForConditionalGeneration.from_pretrained(model_name, trust_remote_code=True)
     nltk.download('punkt')
     print("Model loaded!")
     return tokenizer, model