Update app.py

app.py

@@ -1,50 +1,22 @@
-import streamlit as st
-import torch
-from transformers import pipeline
+import gradio as gr
 
-def main():
-    # Set up the page
-    st.set_page_config(page_title="Nudge Generator Demo - tinyllama 1b", page_icon="orYx logo.png")
-
-    # Title and logo
-    col1, col2 = st.columns([3, 1])
-    with col1:
-        st.title("Nudge Generator Demo - tinyllama 1b")
-    with col2:
-        st.image("orYx logo.png")
-
-    # Chat interface
-    st.markdown("---")
-    st.header("Nudge Response")
-
-    # Input for user-provided message
-    user_message = st.text_area("Enter your message:")
-
-    if st.button("Generate Nudge"):
-        if user_message.strip():
-            with st.spinner("Generating Nudge..."):
-                # Load the pipeline
-                pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="auto")
+# Load the model
+model_interface = gr.load("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
-                # Define the message structure
-                messages = [
-                    {
-                        "role": "system",
-                        "content": "You are a personal corporate trainer that teaches me how to perform well in my organization.",
-                    },
-                    {"role": "user", "content": user_message},
-                ]
-
-                # Generate the prompt
-                prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-                # Generate the response
-                outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-
-                # Display the response
-                st.text_area("Chatbot Response:", outputs[0]["generated_text"], height=200)
-        else:
-            st.warning("Please enter a message to get a response.")
+# Create a Gradio interface with custom title and logo
+def main():
+    with gr.Blocks() as demo:
+        # Add the title
+        gr.Markdown("# Tiny llama 1B Demo")
+        
+        # Add the logo
+        gr.Image("orYx logo.png", elem_id="logo", show_label=False, interactive=False)
+        
+        # Embed the model interface inside a Row
+        with gr.Row():
+            model_interface.render()  # Use render to include the prebuilt interface
+
+    demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()