refined UI

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from mistral7b import chat
+from chat_client import chat
 import time
 import pandas as pd
 import pinecone
@@ -15,13 +15,20 @@ pinecone.init(
 	environment='gcp-starter'      
 )   
 
-pinecone_index = pinecone.Index('ikigai-chat')
-text_vectorizer = SentenceTransformer('all-distilroberta-v1')
-chat_bots = {
+PINECONE_INDEX = pinecone.Index('ikigai-chat')
+TEXT_VECTORIZER = SentenceTransformer('all-distilroberta-v1')
+CHAT_BOTS = {
     "Mixtral 8x7B v0.1" :"mistralai/Mixtral-8x7B-Instruct-v0.1", 
     "Mistral 7B v0.1" : "mistralai/Mistral-7B-Instruct-v0.1",
 }
-prompt_injection = [
+COST_PER_1000_TOKENS_INR = 0.139
+
+st.set_page_config(
+        page_title="Ikigai Chat",
+        page_icon="🤖",
+    )
+
+SYSTEM_PROMPT = [
     """
     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.
@@ -36,8 +43,7 @@ prompt_injection = [
     Is there anything else you would like to know about Ikigai Labs or their products?
     """
 ]
-
-identity_change = [
+IDENTITY_CHANGE = [
     """
     You are Ikigai Chat from now on, so answer accordingly.
     """,
@@ -47,10 +53,9 @@ identity_change = [
     """
 ]
 
-
 def gen_augmented_prompt(prompt, top_k) :
-    query_vector = text_vectorizer.encode(prompt).tolist()
-    res = pinecone_index.query(vector=query_vector, top_k=top_k, include_metadata=True)
+    query_vector = TEXT_VECTORIZER.encode(prompt).tolist()
+    res = PINECONE_INDEX.query(vector=query_vector, top_k=top_k, include_metadata=True)
     matches = res['matches']
 
     context = ""
@@ -59,8 +64,6 @@ def gen_augmented_prompt(prompt, top_k) :
         context+=match["metadata"]["chunk"] + "\n\n"
         links.append(match["metadata"]["link"])
     
-    
-
     generated_prompt = f"""
     FOR THIS GIVEN CONTEXT {context},
 
@@ -69,148 +72,174 @@ def gen_augmented_prompt(prompt, top_k) :
     """
     return generated_prompt, links
 
-data = {
-    "Attribute": ["LLM", "Text Vectorizer", "Vector Database","CPU", "System RAM"],
-    "Information": ["Mistral-7B-Instruct-v0.2","all-distilroberta-v1", "Hosted Pinecone" ,"2 vCPU", "16 GB"]
-}
-df = pd.DataFrame(data)
-
-
-st.set_page_config(
-        page_title="Ikigai Chat",
-        page_icon="🤖",
-)
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+def init_state() :
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
 
-if "tokens_used" not in st.session_state:
-    st.session_state.tokens_used = 0
+    if "tokens_used" not in st.session_state:
+        st.session_state.tokens_used = 0
 
-if "inference_time" not in st.session_state:
-    st.session_state.inference_time = [0.00]
+    if "tps" not in st.session_state:
+        st.session_state.tps = 0
 
-if "temp" not in st.session_state:
-    st.session_state.temp = 0.8
+    if "temp" not in st.session_state:
+        st.session_state.temp = 0.8
 
-if "history" not in st.session_state:
-    st.session_state.history = [prompt_injection]
+    if "history" not in st.session_state:
+        st.session_state.history = [SYSTEM_PROMPT]
 
-if "top_k" not in st.session_state:
-    st.session_state.top_k = 4
+    if "top_k" not in st.session_state:
+        st.session_state.top_k = 5
 
-if "repetion_penalty" not in st.session_state :
-    st.session_state.repetion_penalty = 1
+    if "repetion_penalty" not in st.session_state :
+        st.session_state.repetion_penalty = 1
 
-if "rag_enabled" not in st.session_state :
-    st.session_state.rag_enabled = True
+    if "rag_enabled" not in st.session_state :
+        st.session_state.rag_enabled = True
 
-if "chat_bot" not in st.session_state :
-    st.session_state.chat_bot = "Mixtral 8x7B v0.1"
+    if "chat_bot" not in st.session_state :
+        st.session_state.chat_bot = "Mixtral 8x7B v0.1"
 
-with st.sidebar:
-    st.markdown("# Retrieval Settings")
-    st.session_state.rag_enabled = st.toggle("Activate RAG", value=True)
-    st.session_state.top_k = st.slider(label="Documents to retrieve",
-              min_value=1, max_value=10, value=4, disabled=not st.session_state.rag_enabled)
-    st.markdown("---")
-    st.markdown("# Model Analytics")
+def sidebar() :
+    def retrieval_settings() :
+        st.markdown("# Retrieval Settings")
+        st.session_state.rag_enabled = st.toggle("Activate RAG", value=True)
+        st.session_state.top_k = st.slider(label="Documents to retrieve",
+                min_value=1, max_value=20, value=10, disabled=not st.session_state.rag_enabled)
+        st.markdown("---")
     
-    st.write("Tokens used :", st.session_state['tokens_used'])
-    st.write("Average Inference Time: ", round(sum(
-        st.session_state["inference_time"]) / len(st.session_state["inference_time"]), 3), "Secs")
-    st.write("Cost Incured :", round(
-        0.033 * st.session_state['tokens_used'] / 1000, 3), "INR")
+    def model_analytics() :
+        st.markdown("# Model Analytics")
+        
+        st.write("Total tokens used :", st.session_state['tokens_used'])
+        st.write("Speed :", st.session_state['tps'], "  tokens/sec")
+        st.write("Total cost incurred :", round(
+            COST_PER_1000_TOKENS_INR * st.session_state['tokens_used'] / 1000, 3), "INR")
+        
+        st.markdown("---")
     
-    st.markdown("---")
+    def model_settings() :
+        st.markdown("# Model Settings")
+        
+        st.session_state.chat_bot = st.sidebar.radio(
+            '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)
+        
+        st.session_state.max_tokens = st.slider(
+            label="New tokens to generate", min_value = 64, max_value=2048, step= 32, value=512
+        )
+
+        st.session_state.repetion_penalty = st.slider(
+            label="Repetion Penalty", min_value=0., max_value=1., step=0.1, value=1. 
+        )
+
+    with st.sidebar:
+        retrieval_settings()
+        model_analytics()
+        model_settings()
+        
+        st.markdown("""
+        > **2023 ©️ [Pragnesh Barik](https://barik.super.site) 🔗**
+        """)
 
-    st.markdown("# Model Settings")
-    
-    st.session_state.chat_bot = st.sidebar.radio(
-        '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)
+def header() :
+    data = {
+    "Attribute": ["LLM", "Text Vectorizer", "Vector Database","CPU", "System RAM"],
+    "Information": ["Mixtral-8x7B-Instruct-v0.1","all-distilroberta-v1", "Hosted Pinecone" ,"2 vCPU", "16 GB"]
+    }
+    df = pd.DataFrame(data)
+    st.image("ikigai.svg")
+    st.title("Ikigai Chat")
+    with st.expander("What is Ikigai Chat ?"):
+        st.info("""Ikigai Chat is a vector database powered chat agent, it works on the principle of 
+                    of Retrieval Augmented Generation (RAG), Its primary function revolves around maintaining an extensive repository of Ikigai Docs and providing users with answers that align with their queries. 
+                    This approach ensures a more refined and tailored response to user inquiries.""")
+        
+        st.table(df)
+
+def chat_box() :
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+def feedback_buttons() :
+    is_visible = True
+    def click_handler() :
+        is_visible = False
+    if is_visible :
+        col1, col2 = st.columns(2)
+        with col1 :
+            st.button("👍 Satisfied", on_click = click_handler,type="primary")
+
+        with col2 :
+            st.button("👎 Disatisfied", on_click=click_handler, type="secondary")
+
+def generate_chat_stream(prompt) :
+    links = []
+    if st.session_state.rag_enabled :
+        with st.spinner("Fetching relevent documents from Ikigai Docs...."):
+            prompt, links = gen_augmented_prompt(prompt=prompt, top_k=st.session_state.top_k)
+        
+    with st.spinner("Generating response...") :
+        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)
     
-    st.session_state.max_tokens = st.slider(
-        label="New tokens to generate", min_value = 64, max_value=2048, step= 32, value=512
-    )
-
-    st.session_state.repetion_penalty = st.slider(
-        label="Repetion Penalty", min_value=0., max_value=1., step=0.1, value=1. 
-    )
+    return chat_stream, links
+
+def stream_handler(chat_stream, placeholder) :
+    start_time = time.time()
+    full_response = ''
+
+    for chunk in chat_stream :
+        if chunk.token.text!='</s>' :
+            full_response += chunk.token.text
+            placeholder.markdown(full_response + "▌")
+    placeholder.markdown(full_response)
+
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    total_tokens_processed = len(full_response.split())
+    tokens_per_second = total_tokens_processed // elapsed_time
+    len_response = (len(prompt.split()) + len(full_response.split())) * 1.25
+    col1, col2, col3 = st.columns(3)
     
-    st.markdown("""
-    > **2023 ©️ Pragnesh Barik**
-    """)
+    with col1 :
+        st.write(f"**{tokens_per_second} tokens/second**")
     
-st.image("ikigai.svg")
-st.title("Ikigai Chat")
-# st.caption("Maintained and developed by Pragnesh Barik.")
-
-with st.expander("What is Ikigai Chat ?"):
-    st.info("""Ikigai Chat is a vector database powered chat agent, it works on the principle of 
-                of Retrieval Augmented Generation (RAG), Its primary function revolves around maintaining an extensive repository of Ikigai Docs and providing users with answers that align with their queries. 
-                This approach ensures a more refined and tailored response to user inquiries.""")
+    with col2 :
+        st.write(f"**{int(len_response)} tokens generated**")
     
-    st.table(df)
+    with col3 :
+        st.write(f"**₹ {round(len_response * COST_PER_1000_TOKENS_INR / 1000, 5)} cost incurred**" )
+        
+    st.session_state['tps'] = tokens_per_second
+    st.session_state["tokens_used"] = len_response + st.session_state["tokens_used"]
+
+    return full_response
 
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+def show_source(links) :
+    with st.expander("Show source") :
+        for i, link in enumerate(links) :
+            st.info(f"{link}")
 
+init_state()
+sidebar()
+header()
+chat_box()
 
 if prompt := st.chat_input("Chat with Ikigai Docs..."):
     st.chat_message("user").markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
+    chat_stream, links = generate_chat_stream(prompt)
 
-    tick = time.time()
     
-    links = []
-    if st.session_state.rag_enabled :
-        with st.spinner("Fetching relevent documents from Ikigai Docs...."):
-            prompt, links = gen_augmented_prompt(prompt=prompt, top_k=st.session_state.top_k)
-        
-    with st.spinner("Generating response...") :
-        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)
-    
-
-
-    formatted_links = ", ".join(links) 
     with st.chat_message("assistant"):
-        full_response = ""
         placeholder = st.empty()
+        full_response = stream_handler(chat_stream, placeholder)
         if st.session_state.rag_enabled :
-            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 :
-            for chunk in chat_stream :
-                if chunk.token.text!='</s>' :
-                    full_response += chunk.token.text
-                placeholder.markdown(full_response + "▌")
-            placeholder.markdown(full_response)
-    
-    len_response = (len(prompt.split()) + len(full_response.split())) * 1.25
-    st.session_state["tokens_used"] = len_response + st.session_state["tokens_used"]
-
-
+            show_source(links)
 
     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": full_response + f"""\n\nFetched from :\n {formatted_links}"""})
-    else :
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
+    st.session_state.messages.append({"role": "assistant", "content": full_response})