Update app.py

app.py

@@ -1,33 +1,50 @@
-import gradio as gr
-
-# Load the model from Hugging Face
-def load_model():
-    try:
-        model_interface = gr.load("huggingface/TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-        return model_interface
-    except ValueError as e:
-        print(f"Error loading the model: {e}")
-        return None
-
-# Create a Gradio interface with custom title and logo
+import streamlit as st
+import torch
+from transformers import pipeline
+
 def main():
-    model_interface = load_model()
-    if model_interface is None:
-        print("Failed to load the model. Exiting.")
-        return
-
-    with gr.Blocks() as demo:
-        # Add the title
-        gr.Markdown("# TinyLlama 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()
+    # 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")
+
+                # 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.")
 
 if __name__ == "__main__":
-    main()
+    main()