Update app.py

app.py

@@ -2,11 +2,8 @@ import streamlit as st
 from transformers import pipeline
 
 @st.cache_resource
-def load_model(advanced_mode):
-    if advanced_mode:
-        return pipeline("text-generation", model="PeterBrendan/pbjsGPT2v2")
-    else:
-        return pipeline("text-generation", model="PeterBrendan/pbjs_gpt2")
+def load_model():
+    return pipeline("text-generation", model="PeterBrendan/pbjs_gpt2")
 
 def main():
     if "generated_widget_id" not in st.session_state:
@@ -30,9 +27,6 @@ def main():
     # Create a text input field for custom prompt
     custom_prompt = st.text_input("Enter a custom prompt:", "")
 
-    # Create a checkbox for advanced mode
-    advanced_mode = st.checkbox("Advanced Mode")
-
     # Check if a default prompt is selected
     if default_prompt:
         user_input = default_prompt
@@ -42,7 +36,7 @@ def main():
     # Check if the user input is empty
     if user_input:
         # Load the Hugging Face model
-        generator = load_model(advanced_mode)
+        generator = load_model()
 
         # Display 'Generating Output' message
         output_placeholder = st.empty()
@@ -50,7 +44,7 @@ def main():
             st.write("Generating Output...")
 
         # Generate text based on user input
-        generated_text = generator(user_input, max_length=1100, num_return_sequences=1)[0]["generated_text"]
+        generated_text = generator(user_input, max_length=900, num_return_sequences=1)[0]["generated_text"]
 
         # Clear 'Generating Output' message and display the generated text
         output_placeholder.empty()
@@ -59,4 +53,4 @@ def main():
 
 # Run the app
 if __name__ == "__main__":
-    main()
+    main()