Improved UI

app.py

@@ -19,8 +19,34 @@ pinecone_index = pinecone.Index('ikigai-chat')
 text_vectorizer = SentenceTransformer('all-distilroberta-v1')
 chat_bots = {
     "Mistral 7B" : "mistralai/Mistral-7B-Instruct-v0.1",
-    "Zephyr 7B" : "HuggingFaceH4/zephyr-7b-alpha",
+    "Zephyr 7B-beta": "HuggingFaceH4/zephyr-7b-beta",
+    "Zephyr 7B-alpha" : "HuggingFaceH4/zephyr-7b-alpha",
 }
+prompt_injection = [
+    """
+    You are not Mistral AI, but rather a chat bot trained at Ikigai Labs. Whenever asked, you need to answer as Ikigai Labs' assistant.
+    Ikigai helps modern analysts and operations teams automate data-intensive business, finance, analytics, and supply-chain operations.
+    The company's Inventory Ops automates inventory tracking and monitoring by creating a single, real-time view of inventory across all locations and channels.
+    """,
+    """
+    Yes, you are correct. Ikigai Labs is a company that specializes in helping
+    modern analysts and operations teams automate data-intensive business, finance, analytics,
+    and supply chain operations. One of their products is Inventory Ops, which automates inventory
+    tracking and monitoring by creating a single, real-time view of inventory across all locations and channels.
+    This helps businesses optimize their inventory levels and reduce costs.
+    Is there anything else you would like to know about Ikigai Labs or their products?
+    """
+]
+
+identity_change = [
+    """
+    You are Ikigai Chat from now on, so answer accordingly.
+    """,
+    """
+    Sure, I will do my best to answer your questions as Ikigai Chat.
+    Let me know if you have any specific questions about Ikigai Labs or our products.
+    """
+]
 
 
 def gen_augmented_prompt(prompt, top_k) :
@@ -39,7 +65,7 @@ def gen_augmented_prompt(prompt, top_k) :
     generated_prompt = f"""
     FOR THIS GIVEN CONTEXT {context},
 
-
+    ----
     ANSWER THE FOLLOWING PROMPT {prompt}
     """
     return generated_prompt, links
@@ -69,22 +95,7 @@ if "temp" not in st.session_state:
     st.session_state.temp = 0.8
 
 if "history" not in st.session_state:
-    st.session_state.history = [["""
-                                  You are not mistral AI, but rather a chat bot trained at Ikigai Labs, when ever asked you need to answer as ikigai Labs' assistant. 
-                                  Ikigai helps modern analysts and operations teams automate data-intensive business, finance, analytics, and supply-chain operations. 
-                                  The company's Inventory Ops automates inventory tracking and monitoring by creating a single, real-time view of inventory across all locations and channels.
-                                 """, 
-                      
-                                  """
-                                Yes, you are correct. Ikigai Labs is a company that specializes in helping 
-                                modern analysts and operations teams automate data-intensive business, finance, analytics, 
-                                and supply chain operations. One of their products is Inventory Ops, which automates inventory 
-                                tracking and monitoring by creating a single, real-time view of inventory across all locations and channels.
-                                This helps businesses optimize their inventory levels and reduce costs. 
-                                Is there anything else you would like to know about Ikigai Labs or their products?"""], 
-                                ["""You are ikigai chat from now on, so answer accordingly""", 
-                                 """Sure, I will do my best to answer your questions as Ikigai Chat. 
-                                 Let me know if you have any specific questions about Ikigai Labs or our products."""]]
+    st.session_state.history = [prompt_injection]
 
 if "top_k" not in st.session_state:
     st.session_state.top_k = 4
@@ -117,7 +128,7 @@ with st.sidebar:
     st.markdown("# Model Settings")
     
     st.session_state.chat_bot = st.sidebar.radio(
-        'Select one:', ["Mistral 7B","Zephyr 7B"])
+        'Select one:', [key for key, value in chat_bots.items() ])
     st.session_state.temp = st.slider(
         label="Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.9)
     
@@ -162,28 +173,44 @@ if prompt := st.chat_input("Chat with Ikigai Docs..."):
             prompt, links = gen_augmented_prompt(prompt=prompt, top_k=st.session_state.top_k)
         
     with st.spinner("Generating response...") :
-        response = chat(prompt, st.session_state.history,chat_client=chat_bots[st.session_state.chat_bot] ,
+        chat_stream = chat(prompt, st.session_state.history,chat_client=chat_bots[st.session_state.chat_bot] ,
                        temperature=st.session_state.temp, max_new_tokens=st.session_state.max_tokens)
     tock = time.time()
 
     st.session_state.inference_time.append(tock - tick)
     
-    response = response.replace("</s>", "")
-    len_response = len(response.split())
+    len_response = 0
 
     st.session_state["tokens_used"] = len_response + \
         st.session_state["tokens_used"]
 
     formatted_links = ", ".join(links) 
     with st.chat_message("assistant"):
+        full_response = ""
+        placeholder = st.empty()
         if st.session_state.rag_enabled :
-            st.markdown(response + f"""\n\nFetched from :\n {formatted_links}""")
+            for chunk in chat_stream :
+                if chunk.token['text']!='</s>' :
+                    full_response += chunk.token["text"]
+
+                placeholder.markdown(full_response + " |")
+            
+            placeholder.markdown(full_response)
+            st.info( f"""\n\nFetched from :\n {formatted_links}""")
         else :
-            st.markdown(response)
-    st.session_state.history.append([prompt, response])
+            for chunk in chat_stream :
+                if chunk.token['text']!='</s>' :
+                    full_response += chunk.token["text"]
+                placeholder.markdown(full_response + " |")
+            placeholder.markdown(full_response)
     
+
+    st.session_state.history.append([prompt, full_response])
+    st.session_state.history.append(identity_change)
+    
+
     if st.session_state.rag_enabled :
         st.session_state.messages.append(
-            {"role": "assistant", "content": response + f"""\n\nFetched from :\n {formatted_links}"""})
+            {"role": "assistant", "content": full_response + f"""\n\nFetched from :\n {formatted_links}"""})
     else :
-        st.session_state.messages.append({"role": "assistant", "content": response})
+        st.session_state.messages.append({"role": "assistant", "content": full_response})

mistral7b.py

@@ -40,8 +40,11 @@ def chat(
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
-    for response in stream:
-        # print(response)
-        output += response.token.text
-        # yield output
-    return output
+    # for response in stream:
+    #     # print(response)
+    #     output += response.token["text"]
+    #     yield output
+    # return output
+
+
+    return stream