Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 from langchain_huggingface import HuggingFaceEndpoint
-from transformers import pipeline
+import os
 
 # constants
 QUESTION = "Compute the integral of f(x) = x^2."
@@ -59,12 +59,21 @@ with st.container():
 
             with col1:
                 if st.button("📝 Explain the question"):
+                    prompt = f"[INST]You are a thoughtful AI assistant.\nUser: {QUESTION} [/INST]\nAI:"
+                    st.session_state.response = llm.invoke(prompt)
+                    
                     st.session_state.button_clicked = "Explain the question"
             with col2:
                 if st.button("💡 Give an example"):
+                    prompt = f"[INST]You are a thoughtful AI assistant.\nUser: {QUESTION} [/INST]\nAI:"
+                    st.session_state.response = llm.invoke(prompt)
+                    
                     st.session_state.button_clicked = "Give an example"
             with col3:
                 if st.button("🤔 Who cares?"):
+                    prompt = f"[INST]You are a thoughtful AI assistant.\nUser: {QUESTION} [/INST]\nAI:"
+                    st.session_state.response = llm.invoke(prompt)
+                    
                     st.session_state.button_clicked = "Who cares?"
 
             st.markdown("---")
@@ -72,7 +81,7 @@ with st.container():
         # Display response text if a sub-button is clicked
         if st.session_state.button_clicked:
             with st.container():
-                st.info(f"**hello world  \n{QUESTION}  \n{st.session_state.button_clicked}**")
+                st.info(st.session_state.response)
 
 # Optional: Add footer or spacing
 st.markdown("<br><br>", unsafe_allow_html=True)
@@ -95,7 +104,7 @@ st.markdown(
 )
 
 # source: https://medium.com/@james.irving.phd/creating-your-personal-chatbot-using-hugging-face-spaces-and-streamlit-596a54b9e3ed
-def get_llm_hf_inference(model_id=MODEL, max_new_tokens=130, temperature=0.7):
+def get_llm(model_id=MODEL, max_new_tokens=130, temperature=0.7):
     """
     Returns a language model for HuggingFace inference.