Update main.py

main.py

@@ -4,7 +4,6 @@ from langchain.llms import HuggingFaceHub
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 
-
 class UserInterface():
 
     def __init__(self, ):
@@ -12,11 +11,11 @@ class UserInterface():
         st.text("An Open Source Chat Application")
         st.header("Open LLMs")
 
-        self.API_KEY = st.sidebar.text_input(
-            'API Key',
-            type='password',
-            help="Type in your HuggingFace API key to use this app"
-        )
+        # self.API_KEY = st.sidebar.text_input(
+        #     'API Key',
+        #     type='password',
+        #     help="Type in your HuggingFace API key to use this app"
+        # )
 
         models_name = (
             "HuggingFaceH4/zephyr-7b-beta",
@@ -34,7 +33,7 @@ class UserInterface():
             max_value=1.0,
             step=0.1,
             value=0.5,
-            help="Set the temperature to get the accurate or random result"
+            help="Set the temperature to get accurate or random result"
         )
 
         self.max_token_length = st.sidebar.slider(
@@ -43,7 +42,7 @@ class UserInterface():
             max_value=2048,
             step=16,
             value=64,
-            help="Set max tokens to generate the maximum amount of text output"
+            help="Set max tokens to generate maximum amount of text output"
         )
 
 
@@ -52,7 +51,7 @@ class UserInterface():
             "max_length": self.max_token_length
         }
 
-        os.environ['HUGGINGFACEHUB_API_TOKEN'] = self.API_KEY
+        # os.environ['HUGGINGFACEHUB_API_TOKEN'] = self.API_KEY/
 
     
     def form_data(self):
@@ -61,18 +60,29 @@ class UserInterface():
             if not self.API_KEY.startswith('hf_'):
                 st.warning('Please enter your API key!', icon='⚠')
                 text_input_visibility = True
+            else:
+                text_input_visibility = False
             
+
+            if "messages" not in st.session_state:
+                    st.session_state.messages = []
+
+            st.write(f"You are using {self.models} model")
+
+            for message in st.session_state.messages:
+                with st.chat_message(message.get('role')):
+                    st.write(message.get("content"))
             
-            st.subheader("Context")
-            context = st.text_input(
-                    label="Context",
-                    disabled=text_input_visibility
-                )
-            st.subheader("Question")
-            question = st.text_input(
-                    label="Question",    
-                    disabled=text_input_visibility
-                )
+            context = st.sidebar.text_input(
+                 label="Context",
+                 help="Context lets you know on what the answer should be generated"
+                 )
+
+        
+            question = st.chat_input(
+                 key="question",
+                 disabled=text_input_visibility
+            )
 
 
             template = """
@@ -94,17 +104,41 @@ class UserInterface():
                 model_kwargs = self.model_kwargs
             )
 
-            llm_chain = LLMChain(
-                prompt=prompt,
-                llm=llm,
-            )
+            if question:
+                llm_chain = LLMChain(
+                    prompt=prompt,
+                    llm=llm,
+                )
 
-            result = llm_chain.run({
-                "question": question,
-                "context": context
-            })
+                result = llm_chain.run({
+                    "question": question,
+                    "context": context
+                })
+
+                if "Out of Context" in result:
+                    result = "Out of Context"
+                st.session_state.messages.append(
+                    {
+                        "role":"user",
+                        "content": f"Context: {context}\n\nQuestion: {question}"
+                    }
+                )
+                with st.chat_message("user"):
+                    st.write(f"Context: {context}\n\nQuestion: {question}")
+                
+                if question.lower() == "clear":
+                    del st.session_state.messages
+                    return
+                
+                st.session_state.messages.append(
+                    {
+                        "role": "assistant",
+                        "content": result
+                    }
+                )
+                with st.chat_message('assistant'):
+                    st.markdown(result)
 
-            st.markdown(result)
         except Exception as e:
             st.error(e, icon="🚨")