Update app.py

app.py

@@ -2,8 +2,11 @@ import streamlit as st
 from transformers import pipeline
 
 @st.cache_resource
-def load_model():
-    return pipeline("text-generation", model="PeterBrendan/pbjs_gpt2")
+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 main():
     if "generated_widget_id" not in st.session_state:
@@ -27,6 +30,9 @@ 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
@@ -36,7 +42,7 @@ def main():
     # Check if the user input is empty
     if user_input:
         # Load the Hugging Face model
-        generator = load_model()
+        generator = load_model(advanced_mode)
 
         # Display 'Generating Output' message
         output_placeholder = st.empty()
@@ -44,7 +50,7 @@ def main():
             st.write("Generating Output...")
 
         # Generate text based on user input
-        generated_text = generator(user_input, max_length=700, num_return_sequences=1)[0]["generated_text"]
+        generated_text = generator(user_input, max_length=1500, num_return_sequences=1)[0]["generated_text"]
 
         # Clear 'Generating Output' message and display the generated text
         output_placeholder.empty()