Update app.py

app.py

@@ -1,25 +1,25 @@
 import streamlit as st
 from transformers import pipeline
 
-@st.cache_resource
-def load_model():
-    return pipeline("text-generation", model="PeterBrendan/pbjs_gpt2")
+@st.cache(allow_output_mutation=True)
+def load_model(model_name):
+    return pipeline("text-generation", model=model_name)
 
 def main():
     if "generated_widget_id" not in st.session_state:
         st.session_state["generated_widget_id"] = None
 
     st.title("Prebid Config Generator")
-    st.write("Enter a Prebid config setting, such as 'bidderTimeout', and get a generated Prebid config output starting from that setting onward. Using '{' will generate a Prebid config from the beginning.")
+    st.write("Enter a Prebid config setting, such as 'bidderTimeout', and get a generated Prebid config output starting from that setting onward. Using '{' will generate a Prebid config from the beginning. The model currently has a capped output of 760 characters.")
 
     st.subheader("Intended Uses")
-    st.write("This model is designed to assist publishers in understanding and exploring how other publishers configure their Prebid settings. It can serve as a valuable reference to gain insights into common configurations, best practices, and different approaches used by publishers across various domains.")
-    st.write("To learn more about the model, visit the [pbjs_gpt2 model page](https://huggingface.co/PeterBrendan/pbjs_gpt2). You can also refer to the [official Prebid Documentation on pbjs.setConfig](https://docs.prebid.org/dev-docs/publisher-api-reference/setConfig.html) for more information.")
+    st.write("This model is designed to assist publishers in understanding and exploring how most and advanced publishers configure their Prebid settings. It can serve as a valuable reference to gain insights into common configurations, best practices, and different approaches used by publishers across various domains. The model should be seen as a helpful tool to gain inspiration and understanding of common Prebid settings but not as a substitute for thorough testing and manual review of the final configurations.")
+    st.write("To learn more about the default model, visit the [pbjsGPT2 model page](https://huggingface.co/PeterBrendan/pbjs_gpt2). To learn more about the advanced model, visit the [pbjsGPT2v2 model page](https://huggingface.co/PeterBrendan/pbjsGPT2v2). You can also refer to the [official Prebid Documentation on pbjs.setConfig](https://docs.prebid.org/dev-docs/publisher-api-reference/setConfig.html) for more information.")
 
-    st.write("Note: The model may take some time to generate the output.")
+    st.write("*Note:* The model may take some time to generate the output.")
 
     # Default prompts
-    default_prompts = ["{", "bidderTimeout", "bidderSequence", "Usebidcache", "customPriceBucket"]
+    default_prompts = ["{", "bidderTimeout", "bidderSequence", "Usebidcache", "customPriceBucket", "coppa"]
 
     # Create a selectbox for default prompts
     default_prompt = st.selectbox("Choose a default prompt:", default_prompts)
@@ -35,8 +35,11 @@ def main():
 
     # Check if the user input is empty
     if user_input:
+        # Select the model based on the user's choice
+        model_name = "PeterBrendan/pbjsGPT2v2" if st.button("Advanced Mode") else "PeterBrendan/pbjs_gpt2"
+
         # Load the Hugging Face model
-        generator = load_model()
+        generator = load_model(model_name)
 
         # Display 'Generating Output' message
         output_placeholder = st.empty()
@@ -44,7 +47,7 @@ def main():
             st.write("Generating Output...")
 
         # Generate text based on user input
-        generated_text = generator(user_input, max_length=900, num_return_sequences=1)[0]["generated_text"]
+        generated_text = generator(user_input, max_length=760, num_return_sequences=1)[0]["generated_text"]
 
         # Clear 'Generating Output' message and display the generated text
         output_placeholder.empty()