Update app.py

app.py

@@ -4,18 +4,18 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelWithLMHead
 from transformers import pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("gpt2-large")
+#tokenizer = AutoTokenizer.from_pretrained("gpt2-small")
 
 @st.cache
 def load_model(model_name):
-    model = AutoModelWithLMHead.from_pretrained("gpt2-large")
+    model = AutoModelWithLMHead.from_pretrained(model_name)
     return model
     
 def load_text_gen_model():
-    generator = pipeline("text-generation", model="gpt2-large")
+    generator = pipeline("text-generation", model="gpt2-small")
     return generator 
     
-model = load_model("gpt2-large")
+#model = load_model("gpt2-small")
 
 text_generator = load_text_gen_model()