chatgpt updates, reranking plus copy button plus streaming

app.py

@@ -46,6 +46,15 @@ llmc = Together(
     together_api_key=os.environ['pilotikval']
 )
 
+# Load the reranking model
+reranker = Together(
+    model="mistralai/Mixtral-8x22B-Instruct-v0.1",
+    temperature=0.2,
+    max_tokens=512,
+    top_k=10,
+    together_api_key=os.environ['pilotikval']
+)
+
 msgs = StreamlitChatMessageHistory(key="langchain_messages")
 memory = ConversationBufferMemory(chat_memory=msgs)
 
@@ -63,6 +72,26 @@ def store_chat_history(role: str, content: str):
     # Append the new message to the chat history
     chistory.append({"role": role, "content": content})
 
+def render_message_with_copy_button(role: str, content: str, key: str):
+    html_code = f"""
+    <div class="message" style="position: relative; padding-right: 40px;">
+        <div class="message-content">{content}</div>
+        <button onclick="copyToClipboard('{key}')" style="position: absolute; right: 0; top: 0;">Copy</button>
+    </div>
+    <textarea id="{key}" style="display:none;">{content}</textarea>
+    <script>
+        function copyToClipboard(key) {{
+            var copyText = document.getElementById(key);
+            copyText.style.display = "block";
+            copyText.select();
+            document.execCommand("copy");
+            copyText.style.display = "none";
+            alert("Copied to clipboard");
+        }}
+    </script>
+    """
+    st.write(html_code, unsafe_allow_html=True)
+
 # Define the Streamlit app
 def app():
     with st.sidebar:
@@ -123,9 +152,9 @@ def app():
     conversational_qa_chain = _inputs | _context | ANSWER_PROMPT | llm
 
     st.header("Ask Away!")
-    for message in st.session_state.messages:
+    for i, message in enumerate(st.session_state.messages):
         with st.chat_message(message["role"]):
-            st.write(message["content"])
+            render_message_with_copy_button(message["role"], message["content"], key=f"message-{i}")
             store_chat_history(message["role"], message["content"])
 
     prompts2 = st.chat_input("Say something")
@@ -140,29 +169,37 @@ def app():
             with st.spinner("Thinking..."):
                 for _ in range(3):  # Retry up to 3 times
                     try:
-                        response_generator = stream_conversational_qa_chain(
+                        responses = generate_multiple_responses(
+                            conversational_qa_chain,
                             {
                                 "question": prompts2,
                                 "chat_history": chistory,
-                            }
+                            },
+                            num_responses=5
                         )
-                        message_content = ""
-                        for response_part in response_generator:
-                            message_content += response_part
-                            st.chat_message("assistant", message_content)
-                        st.session_state.messages.append({"role": "assistant", "content": message_content})
+                        best_response = rerank_responses(reranker, responses)
+                        st.write(best_response)
+                        message = {"role": "assistant", "content": best_response}
+                        st.session_state.messages.append(message)
                         break
                     except Exception as e:
                         st.error(f"An error occurred: {e}")
                         time.sleep(2)  # Wait 2 seconds before retrying
 
-def stream_conversational_qa_chain(inputs):
-    try:
-        response = conversational_qa_chain.invoke(inputs)
-        for part in response:
-            yield part
-    except Exception as e:
-        raise e
+def generate_multiple_responses(chain, inputs, num_responses=5):
+    responses = []
+    for _ in range(num_responses):
+        response = chain.invoke(inputs)
+        responses.append(response)
+    return responses
+
+def rerank_responses(reranker, responses):
+    scores = []
+    for response in responses:
+        score = reranker.invoke({"input": response})
+        scores.append(score)
+    best_response_idx = scores.index(max(scores))
+    return responses[best_response_idx]
 
 if __name__ == '__main__':
     app()