Added trigger examples and updated form to implement them, Also reset the session_state.trigger_example to False

app.py

@@ -44,7 +44,7 @@ import torch
 # So first run will load and save resources to global cache, and as user interact and causes rerun of load_model_and_tokenizer(), instead of loading again it will directly use cached resources from memory
 def load_model_and_tokenizer():
     model_name = "google/gemma-2b"  # using gemma-2b for prototype for my GSOC Proposal. Wish me luck.
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name).to("cpu") # As free tier account got CPU only
     # Responsible for automatically downloading and loading the tokenizer configuration and vocabulary associated with the specified pre-trained model.
      # Downloads and loads the tokenizer config and vocab for the given model
     model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -64,7 +64,7 @@ def generate_text(prompt, tone, max_length, temperature=0.7, top_p=0.9, repetiti
     }
     input_text = tone_prompts.get(tone, prompt)
     
-    inputs = tokenizer(input_text, return_tensors="pt")
+    inputs = tokenizer(input_text, return_tensors="pt").to("cpu") # Don't make me say again, Free tier.
     outputs = model.generate(
         inputs["input_ids"],
         max_length=max_length + len(input_text.split()),
@@ -110,25 +110,28 @@ with st.expander("\U0001F9E0 How does this work? Click to peek inside."):
     """)
 
 # One-click examples
-example_clicked = None
+if "trigger_example" not in st.session_state:
+    st.session_state.trigger_example = False
+
 col1, col2 = st.columns(2)
 with col1:
     if st.button("Try Funny Cat Story"):
-        prompt = "The cat hacked my WiFi"
-        tone = "Funny"
-        example_clicked = True
+        st.session_state.prompt = "The cat hacked my WiFi"
+        st.session_state.tone = "Funny"
+        st.session_state.trigger_example = True
 with col2:
     if st.button("Try Poetic Goodbye"):
-        prompt = "As the sun set on our final day"
-        tone = "Poetic"
-        example_clicked = True
+        st.session_state.prompt = "As the sun set on our final day"
+        st.session_state.tone = "Poetic"
+        st.session_state.trigger_example = True
+
 
 
 
 # User input section
 with st.form(key="input_form"):
-    prompt = st.text_input("Enter a prompt", placeholder="e.g., 'The future of AI is'")
-    tone = st.selectbox("Tone", ["Funny", "Serious", "Poetic"])
+    prompt = st.text_input("Enter a prompt", placeholder="e.g., 'The future of AI is'", value=st.session_state.get("prompt", ""))
+    tone = st.selectbox("Tone", ["Funny", "Serious", "Poetic"], index=["Funny", "Serious", "Poetic"].index(st.session_state.get("tone", "Funny")))
     temperature = st.slider("Temperature (Creativity)", 0.2, 1.5, 0.7)
     top_p = st.slider("Top-p (Nucleus Sampling)", 0.1, 1.0, 0.9)
     repetition_penalty = st.slider("Repetition Penalty", 1.0, 2.0, 1.0)
@@ -137,7 +140,8 @@ with st.form(key="input_form"):
     submit_button = st.form_submit_button(label="Generate")
 
 # Generate and display output
-if submit_button:
+if submit_button or st.session_state.trigger_example:
+    st.session_state.trigger_example = False  # Reset after use
     if not prompt:
         st.error("Please enter a prompt!")
     else: