Updated import

app.py

@@ -3,7 +3,7 @@
 """
 
 # Install necessary libraries
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, GPT2LMHeadModel, pipeline
 import streamlit as st
 from pprint import pprint
 import json
@@ -19,7 +19,7 @@ st.set_page_config(layout="wide")
 @st.cache(allow_output_mutation=True)
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(cfg["model_name_or_path"])
-    model = AutoModelForCausalLM.from_pretrained(cfg["model_name_or_path"])
+    model = GPT2LMHeadModel.from_pretrained(cfg["model_name_or_path"])
     generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
     return generator, tokenizer