Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
 
 def main():
     # Set up the page
@@ -19,21 +19,11 @@ def main():
     # Input for user-provided data
     prompt = st.text_area("Enter the prompt here:")
 
-    # Load model directly
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
-    model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it")
-
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-    # Button to generate the nudge
-    if st.button("Generate Nudge"):
-        if prompt.strip():
-            with st.spinner("Generating nudges..."):
-                response = pipe(messages)
-                st.success("Nudges generated successfully!")
-                st.text_area("Pipeline Response:", str(response), height=200)
-        else:
-            st.warning("Please enter data to generate nudges.")
+    # Load model with Gradio
+    if st.button("Load and Launch Model"):
+        with st.spinner("Launching model..."):
+            gr.load("models/google/gemma-2-2b-it").launch()
+            st.success("Model launched successfully! Access it in the new tab opened by Gradio.")
 
 if __name__ == "__main__":
     main()