import streamlit as st
import gradio as gr

def main():
    # Set up the page
    st.set_page_config(page_title="Nudge Generator - Gemma 2b", page_icon="orYx logo.png", layout="wide")

    # Title and logo
    col1, col2 = st.columns([3, 1])
    with col1:
        st.title("Nudge Generator - Gemma 2b")
    with col2:
        st.image("orYx logo.png", use_column_width=True)

    # Chat interface
    st.markdown("---")
    st.header("Chat Interface")

    # Input for user-provided data
    prompt = st.text_area("Enter the prompt here:")

    # 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()