Upload 2 files

app.py

@@ -4,6 +4,9 @@ from streamlit.components.v1 import html
 from streamlit_extras.stylable_container import stylable_container
 import re
 import urllib.parse
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+import numpy as np
+import pypdfium2 as pdfium
 
 st.title("Cohere Chat UI")
 
@@ -12,7 +15,7 @@ if "api_key" not in st.session_state:
     if api_key:
         if api_key.isascii():
             st.session_state.api_key = api_key
-            client = cohere.Client(api_key=api_key)
+            client = cohere.ClientV2(api_key=api_key)
             st.rerun()
         else:
             st.warning("Please enter your API key correctly.")
@@ -21,30 +24,84 @@ if "api_key" not in st.session_state:
         st.warning("Please enter your API key to use the app. You can obtain your API key from here: https://dashboard.cohere.com/api-keys")
         st.stop()
 else:
-    client = cohere.Client(api_key=st.session_state.api_key)
+    client = cohere.ClientV2(api_key=st.session_state.api_key)
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-def get_ai_response(prompt, chat_history):
+if "rag_file_key" not in st.session_state:
+    st.session_state.rag_file_key = None
+
+if "rag_embedded" not in st.session_state:
+    st.session_state.rag_embedded = False
+
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=512,
+    chunk_overlap=50,
+    length_function=len,
+    is_separator_regex=False,
+)
+
+def batch_embed(texts, batch_size=96):
+    all_embeddings = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i:i+batch_size]
+        response = client.embed(
+            texts=batch,
+            model=embed_model,
+            input_type="search_document",
+            embedding_types=['float']
+        )
+        all_embeddings.extend(response.embeddings.float)
+    return all_embeddings
+
+def cosine_similarity(a, b):
+    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
+
+def get_ai_response(chat_history):
     st.session_state.is_streaming = True
     st.session_state.response = ""
     
-    with st.chat_message("ai", avatar=st.session_state.assistant_avatar):
+    with st.chat_message("assistant", avatar=st.session_state.assistant_avatar):
+        # RAG
+        if st.session_state.get("rag_chunks") and st.session_state.get("rag_embeddings"):
+            chunks = st.session_state.rag_chunks
+            embeddings = st.session_state.rag_embeddings
+
+            vector_database = {i: np.array(embedding) for i, embedding in enumerate(embeddings)}
+
+            query = chat_history[-1]["content"]
+            query_embedding = client.embed(texts=[query], model=embed_model, input_type="search_query", embedding_types=['float']).embeddings.float[0]
+
+            similarities = [cosine_similarity(query_embedding, chunk) for chunk in embeddings]
+            top_indices = np.argsort(similarities)[::-1][:10]
+            top_chunks_after_retrieval = [chunks[i] for i in top_indices]
+
+            rerank_response = client.rerank(query=query, documents=top_chunks_after_retrieval, top_n=3, model=rerank_model)
+            top_chunks_after_rerank = [top_chunks_after_retrieval[result.index] for result in rerank_response.results]
+            documents = [{"data": {"title": f"chunk {i}", "snippet": chunk}} for i, chunk in enumerate(top_chunks_after_rerank)]
+        
         penalty_kwargs = {
             "frequency_penalty" if penalty_type == "Frequency Penalty" else "presence_penalty": penalty_value
         }
-        
-        stream = client.chat_stream(
-            message=prompt,
-            model=model,
-            preamble=preamble,
-            chat_history=chat_history,
-            temperature=temperature,
-            k=k,
-            p=p,
+
+        chat_history.insert(0, {"role": "system", "content": preamble})
+
+        stream_kwargs = {
+            "messages": chat_history,
+            "model": model,
+            "temperature": temperature,
+            "k": k,
+            "p": p,
             **penalty_kwargs
-        )
+        }
+        
+        if st.session_state.get("rag_text"):
+            stream_kwargs["documents"] = documents
+        elif model in ["command-r-08-2024", "command-r-plus-08-2024"]:
+            stream_kwargs["safety_mode"] = "OFF"
+        
+        stream = client.chat_stream(**stream_kwargs)
 
         placeholder = st.empty()
 
@@ -63,10 +120,10 @@ def get_ai_response(prompt, chat_history):
             st.button("Stop generating")
 
         shown_message = ""
-        
-        for event in stream:
-            if event.event_type == "text-generation":
-                content = event.text
+
+        for chunk in stream:
+            if chunk.type == "content-delta":
+                content = chunk.delta.message.content.text
                 st.session_state.response += content
                 shown_message += content.replace("\n", "  \n")\
                                         .replace("<", "\\<")\
@@ -90,10 +147,9 @@ inline_pattern = r"`([^`\n]+?)`"
 
 def display_messages():
     for i, message in enumerate(st.session_state.messages):
-        name = "user" if message["role"] == "USER" else "ai"
-        avatar = st.session_state.user_avatar if message["role"] == "USER" else st.session_state.assistant_avatar
-        with st.chat_message(name, avatar=avatar):
-            shown_message = message["text"].replace("\n", "  \n")\
+        avatar = st.session_state.user_avatar if message["role"] == "user" else st.session_state.assistant_avatar
+        with st.chat_message(message["role"], avatar=avatar):
+            shown_message = message["content"].replace("\n", "  \n")\
                                            .replace("<", "\\<")\
                                            .replace(">", "\\>")
             if "```" in shown_message:
@@ -113,7 +169,7 @@ def display_messages():
                     del st.session_state.messages[i]
                     st.rerun()
             with col3:
-                text_to_copy = message["text"]
+                text_to_copy = message["content"]
                 # Encode the string to escape
                 text_to_copy_escaped = urllib.parse.quote(text_to_copy)
 
@@ -131,7 +187,7 @@ def display_messages():
                 """
                 html(copy_button_html, height=50)
 
-            if i == len(st.session_state.messages) - 1 and message["role"] == "CHATBOT":
+            if i == len(st.session_state.messages) - 1 and message["role"] == "assistant":
                 with col4:
                     if st.button("Retry", key=f"retry_{i}_{len(st.session_state.messages)}"):
                         if len(st.session_state.messages) >= 2:
@@ -141,11 +197,11 @@ def display_messages():
         
         if "edit_index" in st.session_state and st.session_state.edit_index == i:
             with st.form(key=f"edit_form_{i}_{len(st.session_state.messages)}"):
-                new_content = st.text_area("Edit message", height=200, value=st.session_state.messages[i]["text"])
+                new_content = st.text_area("Edit message", height=200, value=st.session_state.messages[i]["content"])
                 col1, col2 = st.columns([1, 1])
                 with col1:
                     if st.form_submit_button("Save"):
-                        st.session_state.messages[i]["text"] = new_content
+                        st.session_state.messages[i]["content"] = new_content
                         del st.session_state.edit_index
                         st.rerun()
                 with col2:
@@ -163,12 +219,12 @@ with st.sidebar:
         # Copy Conversation History button
         log_text = ""
         for message in st.session_state.messages:
-            if message["role"] == "USER":
+            if message["role"] == "user":
                 log_text += "<USER>\n"
-                log_text += message["text"] + "\n\n"
+                log_text += message["content"] + "\n\n"
             else:
                 log_text += "<ASSISTANT>\n"
-                log_text += message["text"] + "\n\n"
+                log_text += message["content"] + "\n\n"
         log_text = log_text.rstrip("\n")
             
         # Encode the string to escape
@@ -178,12 +234,21 @@ with st.sidebar:
         <button id="copy-log-btn" style='font-size: 1em; padding: 0.5em;' onclick='copyLog()'>Copy Conversation History</button>
         
         <script>
+        const log_text_escaped = "{log_text_escaped}";
         function copyLog() {{
-            navigator.clipboard.writeText(decodeURIComponent("{log_text_escaped}"));
-            let copyBtn = document.getElementById("copy-log-btn");
+            navigator.clipboard.writeText(decodeURIComponent(log_text_escaped));
+            const copyBtn = document.getElementById("copy-log-btn");
             copyBtn.innerHTML = "Copied!";
             setTimeout(function(){{ copyBtn.innerHTML = "Copy Conversation History"; }}, 2000);
         }}
+        window.parent.document.addEventListener('keydown', (e) => {{
+            if ( e.code == "Pause" ){{
+                window.parent.navigator.clipboard.writeText(decodeURIComponent(log_text_escaped));
+                const copyBtn = document.getElementById("copy-log-btn");
+                copyBtn.innerHTML = "Copied!";
+                setTimeout(function(){{ copyBtn.innerHTML = "Copy Conversation History"; }}, 2000);
+            }}
+        }} , false);
         </script>
         """
         html(copy_log_button_html, height=50)
@@ -208,7 +273,56 @@ with st.sidebar:
         p = st.slider("Top-P", min_value=0.01, max_value=0.99, value=0.75, step=0.01)
         penalty_type = st.selectbox("Penalty Type", options=["Frequency Penalty", "Presence Penalty"])
         penalty_value = st.slider("Penalty Value", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
-    
+
+        st.header("RAG")
+        st.markdown("Select the model and encoding before uploading the file.")
+        rag_model = st.selectbox("RAG Model", options=["Multilingual", "English"], index=0)
+        file_encoding = st.selectbox("Encoding", options=["utf_8", "shift_jis"], index=0)
+        st.session_state.rag_file = st.file_uploader("Choose a txt or pdf file", type=["txt", "pdf"], key="rag_file_uploader")
+
+        if rag_model == "Multilingual":
+            embed_model = "embed-multilingual-v3.0"
+            rerank_model = "rerank-multilingual-v3.0"
+        else:
+            embed_model = "embed-english-v3.0"
+            rerank_model = "rerank-english-v3.0"
+        
+        if st.session_state.rag_file is not None:
+            if st.session_state.rag_file_key != st.session_state.rag_file:
+                st.session_state.rag_file_key = st.session_state.rag_file
+                st.session_state.rag_embedded = False
+                if "rag_text" in st.session_state:
+                    del st.session_state.rag_text
+                if "rag_chunks" in st.session_state:
+                    del st.session_state.rag_chunks
+                if "rag_embeddings" in st.session_state:
+                    del st.session_state.rag_embeddings
+        
+            if not st.session_state.rag_embedded:
+                if st.session_state.rag_file.type == "application/pdf":
+                    pdf = pdfium.PdfDocument(st.session_state.rag_file)
+                    st.session_state.rag_text = ""
+                    for page in pdf:
+                        textpage = page.get_textpage()
+                        st.session_state.rag_text += textpage.get_text_range()
+                else:
+                    st.session_state.rag_text = st.session_state.rag_file.read().decode(file_encoding)
+                chunks_ = text_splitter.create_documents([st.session_state.rag_text])
+                chunks = [c.page_content for c in chunks_]
+                embeddings = batch_embed(chunks)
+                st.session_state.rag_chunks = chunks
+                st.session_state.rag_embeddings = embeddings
+                st.session_state.rag_embedded = True
+        else:
+            st.session_state.rag_file_key = None
+            st.session_state.rag_embedded = False
+            if "rag_text" in st.session_state:
+                del st.session_state.rag_text
+            if "rag_chunks" in st.session_state:
+                del st.session_state.rag_chunks
+            if "rag_embeddings" in st.session_state:
+                del st.session_state.rag_embeddings
+        
         st.header("Restore History")
         history_input = st.text_area("Paste conversation history:", height=200)
         if st.button("Restore History"):
@@ -219,13 +333,13 @@ with st.sidebar:
             for message in messages:
                 if message.strip() in ["<USER>", "<ASSISTANT>"]:
                     if role and text:
-                        st.session_state.messages.append({"role": role, "text": text.strip()})
+                        st.session_state.messages.append({"role": role, "content": text.strip()})
                         text = ""
-                    role = "USER" if message.strip() == "<USER>" else "CHATBOT"
+                    role = "user" if message.strip() == "<USER>" else "assistant"
                 else:
                     text += message
             if role and text:
-                st.session_state.messages.append({"role": role, "text": text.strip()})
+                st.session_state.messages.append({"role": role, "content": text.strip()})
             st.rerun()
     
         st.header("Clear History")
@@ -238,13 +352,14 @@ with st.sidebar:
         if st.button("Update API Key"):
             if new_api_key and new_api_key.isascii():
                 st.session_state.api_key = new_api_key
-                client = cohere.Client(api_key=new_api_key)
+                client = cohere.ClientV2(api_key=new_api_key)
                 st.success("API Key updated successfully!")
             else:
                 st.warning("Please enter a valid API Key.")
                 
     with appearance_tab:
         st.header("Font Selection")
+        
         font_options = {
             "Zen Maru Gothic": "Zen Maru Gothic",
             "Noto Sans JP": "Noto Sans JP",
@@ -267,7 +382,7 @@ with st.sidebar:
 
 # After Stop generating
 if st.session_state.get("is_streaming"):
-    st.session_state.messages.append({"role": "CHATBOT", "text": st.session_state.response})
+    st.session_state.messages.append({"role": "assistant", "content": st.session_state.response})
     st.session_state.is_streaming = False
     if "retry_flag" in st.session_state and st.session_state.retry_flag:
         st.session_state.retry_flag = False
@@ -311,16 +426,16 @@ display_messages()
 # After Retry
 if st.session_state.get("retry_flag"):
     if len(st.session_state.messages) > 0:
-        prompt = st.session_state.messages[-1]["text"]
-        messages = st.session_state.messages[:-1].copy()
-        response = get_ai_response(prompt, messages)
-        st.session_state.messages.append({"role": "CHATBOT", "text": response})
+        messages = st.session_state.messages.copy()
+        response = get_ai_response(messages)
+        st.session_state.messages.append({"role": "assistant", "content": response})
         st.session_state.retry_flag = False
         st.rerun()
     else:
         st.session_state.retry_flag = False
 
 if prompt := st.chat_input("Enter your message here..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
     chat_history = st.session_state.messages.copy()
     
     shown_message = prompt.replace("\n", "  \n")\
@@ -330,9 +445,7 @@ if prompt := st.chat_input("Enter your message here..."):
     with st.chat_message("user", avatar=st.session_state.user_avatar):
         st.write(shown_message)
     
-    st.session_state.messages.append({"role": "USER", "text": prompt})
-    
-    response = get_ai_response(prompt, chat_history)
+    response = get_ai_response(chat_history)
     
-    st.session_state.messages.append({"role": "CHATBOT", "text": response})
+    st.session_state.messages.append({"role": "assistant", "content": response})
     st.rerun()

requirements.txt

@@ -1,2 +1,5 @@
 cohere
-streamlit-extras
+streamlit-extras
+numpy
+langchain-text-splitters
+pypdfium2