Update app.py

app.py

@@ -40,14 +40,14 @@
 # if __name__ == "__main__":
 #     main()
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 # Load the DialoGPT model and tokenizer
 @st.cache_resource
 def load_model():
     model_name = "microsoft/DialoGPT-medium"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+    model = GPT2LMHeadModel.from_pretrained(model_name)
     return model, tokenizer
 
 # Function to generate a response from DialoGPT
@@ -80,3 +80,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+