Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,6 @@
 import os
 os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit"
-from langchain_huggingface import ChatHuggingFace
+from langchain_huggingface import HuggingFaceEndpoint
 import streamlit as st
 
 # constants
@@ -19,10 +19,12 @@ if "response" not in st.session_state:
     st.session_state.response = ""
 
 def get_llm(model_id=MODEL, max_new_tokens=300, temperature=0.7):
-    return ChatHuggingFace(
+    os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")  # Optional but ensures it's set
+
+    return HuggingFaceEndpoint(
         repo_id=model_id,
-        temperature=temperature,
         max_new_tokens=max_new_tokens,
+        temperature=temperature,
     )
     
 # create llm
@@ -54,6 +56,11 @@ p_application = """
     6. Do not include any explanation or examples in your response.
 """
 
+# count the number of times "I don't know is clicked"
+if "retry_count" not in st.session_state:
+    st.session_state.retry_count = 0
+
+    
 # Initialize session state
 if "help_clicks" not in st.session_state:
     st.session_state.help_clicks = 0
@@ -149,19 +156,16 @@ with st.container():
         if st.session_state.button_clicked:
             with st.container():
                 st.info(st.session_state.response)
-
-                if st.session_state.button_clicked:
-                    with st.container():
-                        st.info(st.session_state.response)
                 
-                        if st.session_state.button_clicked == "Explain the question":
-                            if st.button("I don't understand. Try again."):
-                                st.session_state.retry_count += 1
-                                alt_llm = get_llm(temperature=0.9)
-                                st.session_state.response = alt_llm.invoke(
-                                    st.session_state.full_prompt
-                                )
-
+                if st.session_state.button_clicked == "Explain the question":
+                    col1, col2, col3 = st.columns([1, 1, 1])
+                    with col2:
+                        if st.button("I don't understand. Try again.", key="retry_button"):
+                            st.session_state.retry_count += 1
+                            alt_llm = get_llm(temperature=0.9)
+                            st.session_state.response = alt_llm.invoke(
+                                st.session_state.full_prompt
+                            )
 # Optional: Add footer or spacing
 st.markdown("<br><br>", unsafe_allow_html=True)