Update app.py

app.py

@@ -4,6 +4,8 @@ from sklearn.metrics.pairwise import cosine_similarity
 import re
 from huggingface_hub import InferenceClient
 import os
+import numpy as np
+
 
 
 
@@ -17,7 +19,6 @@ my_initial_rag_text = f"""This is a RAG (Retrieval-Augmented Generation) chatbot
 - Uses Streamlit for the web interface
 - Employs SentenceTransformer for generating embeddings
 - Uses HuggingFace's InferenceClient for LLM interaction
-- Has a default text about a GRNET training module on LLMs
 
 2. State Management:
 - Maintains several session state variables for:
@@ -83,24 +84,39 @@ if "embeddings_model" not in st.session_state:
     # We will use the all-MiniLM-L6-v2 model for embeddings
     st.session_state["embeddings_model"] = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
-MAXIMUM_TOKENS = 512
 
 my_system_instructions = "You are a helpful assistant. Be brief and concise. Provide your answers in 100 words or less."
 
 first_message = "Hello, how can I help you today?"
 
-# Check if the chat messages are not already in the session state
-if "my_chat_messages" not in st.session_state:
-    # Initialize the chat messages list in the session state
-    st.session_state["my_chat_messages"] = []
-    # Add the system instructions to the chat messages
-    st.session_state["my_chat_messages"].append({"role": "system", "content": my_system_instructions})
-
 def delete_chat_messages():
     for key in st.session_state.keys():
-        if key != "my_rag_text":
+        if key != "my_rag_text" and key != "my_system_instructions":
             del st.session_state[key]
 
+def create_sentences_rag():
+    text = st.session_state["my_rag_text"]
+
+    # Split only on .?!;: followed by space OR on \n
+    pattern = r'(?<=[.?!;:])\s+|\n'
+    st.session_state["my_sentences"] = [sentence.strip() for sentence in re.split(pattern, text) if sentence.strip()]
+    sentences_ids = [i for i in range(len(st.session_state["my_sentences"]))]
+
+    # Rolling window: include partial windows at end
+    st.session_state["my_sentences_rag_ids"] = []
+    st.session_state["my_sentences_rag"] = []
+    for rolling_window_size in range(st.session_state["min_window_size"], st.session_state["max_window_size"]+1):
+        for i in range(0, len(st.session_state["my_sentences"])-rolling_window_size+1):
+            chunk = " ".join(st.session_state["my_sentences"][i:i+rolling_window_size]).strip()
+            if chunk:   
+                st.session_state["my_sentences_rag"].append(chunk)
+                st.session_state["my_sentences_rag_ids"].append(sentences_ids[i:i+rolling_window_size])
+                # print(f"*****{chunk}*****\n")
+
+    print(len(st.session_state["my_sentences_rag"]))
+    st.session_state["my_embeddings"] = st.session_state["embeddings_model"].encode(st.session_state["my_sentences_rag"])
+
+
 augmented_prompt = ""
 
 # Create two columns with a 1:2 ratio
@@ -118,7 +134,7 @@ Large Language Models may provide wrong answers. Please verify the answers and c
 The user agrees to indemnify and hold harmless the developers of the Software from any related claims or disputes arising from the utilization of the Software by the user.
 
 By using the Software, you agree to the terms and conditions of the disclaimer.""")
-    
+
     # Add a selectbox for model selection
     st.selectbox("Select the model to use:",
                 ["mistralai/Mistral-7B-Instruct-v0.3",
@@ -126,44 +142,46 @@ By using the Software, you agree to the terms and conditions of the disclaimer."
                 "HuggingFaceH4/zephyr-7b-beta"], 
                 key="my_llm_model", on_change=update_llm_model)
     
+    # Add a text are for the system instructions
+    st.text_area(label="Please enter your system instructions here:", value=my_system_instructions, height=100, key="my_system_instructions", on_change=delete_chat_messages)
+    
     # Add a text area for RAG text input
     st.text_area(label="Please enter your RAG text here:", value=my_initial_rag_text, height=500, key="my_rag_text", on_change=delete_chat_messages)
 
+    # Add a slider for minimum window size
+    st.slider("Minimum window size in original sentences", min_value=1, max_value=20, value=5, step=1, key="min_window_size", on_change=create_sentences_rag)
+
+    # Add a slider for maximum window size
+    st.slider("Maximum window size in original sentences", min_value=1, max_value=20, value=10, step=1, key="max_window_size", on_change=create_sentences_rag)
+    
+    # Add a slider for the similarity threshold
+    st.slider("Similarity threshold", min_value=0.0, max_value=1.0, value=0.3, step=0.01, key="my_similarity_threshold")
 
+    # Add a slider for the number of sentences to keep
+    st.slider("Number of original chunks to keep", min_value=1, max_value=50, value=20, step=1, key="nof_keep_sentences")
 
+    # Add a slider for the number of minimum sub prompts
+    st.slider("Minimum number of words in sub prompt split", min_value=1, max_value=10, value=1, step=1, key="nof_min_sub_prompts")
 
-# Define parameters
-min_window_size = 5 # number of sentences per chunk
-max_window_size = 10 # number of sentences per chunk
-print(100*"-")
-# Check if the sentences are not already in the session state
-if "my_sentences" not in st.session_state:
-    text = st.session_state["my_rag_text"]
+    # Add a slider for the number of maximum sub prompts
+    st.slider("Maximum number of words in sub prompt split", min_value=1, max_value=10, value=5, step=1, key="nof_max_sub_prompts")
 
-    # Split only on .?!;: followed by space OR on \n
-    pattern = r'(?<=[.?!;:])\s+|\n'
-    sentence_split = re.split(pattern, text)
 
-    sentences = [s.strip() for s in sentence_split if s.strip()]
 
-    # Rolling window: include partial windows at end
-    chunks = []
-    for rolling_window_size in range(min_window_size, max_window_size+1):
-        for i in range(0, len(sentences)-rolling_window_size+1):
-            chunk = " ".join(sentences[i:i+rolling_window_size]).strip()
-            if chunk:   
-                chunks.append(chunk)
-                print(f"*****{chunk}*****\n")
+# Check if the chat messages are not already in the session state
+if "my_chat_messages" not in st.session_state:
+    # Initialize the chat messages list in the session state
+    st.session_state["my_chat_messages"] = []
+    # Add the system instructions to the chat messages
+    st.session_state["my_chat_messages"].append({"role": "system", "content": st.session_state["my_system_instructions"]})
 
-    st.session_state["my_sentences"] = chunks
-    print(len(st.session_state["my_sentences"]))
 
+# print(100*"-")
+# Check if the sentences are not already in the session state
+if "my_sentences_rag" not in st.session_state:
+    create_sentences_rag()
 
 
-# Check if the embeddings are not already in the session state
-if "my_embeddings" not in st.session_state:
-    st.session_state["my_embeddings"] = st.session_state["embeddings_model"].encode(st.session_state["my_sentences"])
-    
 with column_2:
     # Create a container for the messages with a specified height
     messages_container = st.container(height=500)
@@ -183,29 +201,57 @@ with column_2:
     # Check if there is a new prompt from the user
     if prompt := st.chat_input("you may ask here your questions"):
         
-        # Encode the user's prompt to get its embedding
-        my_question_embedding = st.session_state.embeddings_model.encode([prompt])
+        split_prompt = prompt.split(" ")
+        all_sub_prompts = []
+        for jj in range(st.session_state["nof_min_sub_prompts"], st.session_state["nof_max_sub_prompts"]+1):
+            for ii in range(len(split_prompt)):
+                i_split = " ".join(split_prompt[ii:ii+jj]).strip()
+                if i_split:
+                    all_sub_prompts.append(i_split)
         
-        # Calculate the cosine similarity between the prompt embedding and stored embeddings
-        similarity_to_question = cosine_similarity(my_question_embedding, st.session_state.my_embeddings).flatten()
-        
-        # Number of sentences to keep based on similarity
-        nof_keep_sentences = 3
+        similarities_to_question = np.zeros(len(st.session_state["my_embeddings"]))
+        for sub_prompt in all_sub_prompts:
+            # Encode the user's prompt to get its embedding
+            my_question_embedding = st.session_state.embeddings_model.encode([sub_prompt])
+            
+            # Calculate the cosine similarity between the prompt embedding and stored embeddings
+            similarities_to_question += cosine_similarity(my_question_embedding, st.session_state["my_embeddings"]).flatten()
+        similarities_to_question /= len(all_sub_prompts)
         
         # Get the indices of the top similar sentences
-        sorted_indices = similarity_to_question.argsort()[::-1][:nof_keep_sentences][::-1]
-        
-        # Retrieve the top similar sentences
-        sorted_sentences = ["Importance: " + str(round(100*similarity_to_question[i])) + f"% {5*'>'} " + st.session_state.my_sentences[i] for i in sorted_indices]
-        
+        bottom_col1, bottom_col2 = st.columns([1, 1])
+        sorted_indices_rag = similarities_to_question.argsort()[::-1]
+        sorted_indices_sentences = []
+        max_similarity = 0
+        # for irag in range(st.session_state["nof_keep_sentences"]):
+        irag = 0
+        while len(set(sorted_indices_sentences))<st.session_state["nof_keep_sentences"]:
+            sorted_indices_sentences.extend(st.session_state["my_sentences_rag_ids"][sorted_indices_rag[irag]])
+            max_similarity = max(max_similarity, similarities_to_question[sorted_indices_rag[irag]])
+            with bottom_col1:
+                str_conf = f"Confidence: {similarities_to_question[sorted_indices_rag[irag]]:.5f}, Sentences IDs: {st.session_state["my_sentences_rag_ids"][sorted_indices_rag[irag]]}"
+                with st.expander(f"Chunk: {str(irag+1)} {str_conf}"):
+                    for idx in st.session_state["my_sentences_rag_ids"][sorted_indices_rag[irag]]:
+                        st.write(f"{st.session_state["my_sentences"][idx]}")
+            irag += 1
+
+        sorted_indices_sentences = sorted(list(set(sorted_indices_sentences)))
+
         # Construct the augmented prompt with the similar sentences
-        augmented_prompt = "Here is my context:"
-        for sentence in sorted_sentences:
-            augmented_prompt += "\n\n" + 20*"-" + f"\n\n{sentence}"
-        augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "use the context I gave you to reply on the following:"
-        augmented_prompt += "\n\n" + f"\n\n{prompt}"
+        if max_similarity > st.session_state["my_similarity_threshold"]:
+            augmented_prompt = "This is my context:" + "\n\n" + 20*"-" + "\n\n" 
+            augmented_prompt += "\n".join([st.session_state["my_sentences"][idx] for idx in sorted_indices_sentences])
+            augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "If the above context is not relevant to the prompt, ignore the context and reply based only on the prompt."
+            augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "If the above context is relevant to the prompt, reply based on the context and the prompt."
+            augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "The prompt is:"
+            augmented_prompt += "\n\n" + f"\n\n{prompt}"
+        else:
+            augmented_prompt = prompt
+        with bottom_col2:
+            # Display the augmented prompt used for generating the response
+            st.write("Augmented prompt:")
+            st.json({"max_similarity": max_similarity, "augmented_prompt": augmented_prompt}, expanded=False)
         
-
         # Display the user's prompt in the chat container with a specific avatar
         messages_container.chat_message("user", avatar=":material/psychology_alt:").markdown(prompt) 
         # Append the augmented prompt to the chat messages in the session state
@@ -228,14 +274,9 @@ with column_2:
         # Append the assistant's response to the chat messages in the session state
         st.session_state["my_chat_messages"].append({"role": "assistant", "content": response})
 
-    
-    # Display the chat messages history
-    st.write("Messages History All:")
-    st.json(st.session_state["my_chat_messages"], expanded=False)
-    # Display the augmented prompt used for generating the response
-    st.write("Augmented prompt:")
-    st.json({"augmented_prompt": augmented_prompt}, expanded=False)
+        with bottom_col2:
+            # Display the chat messages history
+            st.write("Messages History All:")
+            st.json(st.session_state["my_chat_messages"], expanded=False)
 
-    
 
-