Spaces:
Sleeping
Sleeping
| import openai | |
| import streamlit as st | |
| from streamlit.components.v1 import html | |
| from streamlit_extras.stylable_container import stylable_container | |
| import re | |
| import urllib.parse | |
| import traceback | |
| import tiktoken | |
| st.title("Nvidia Chat UI") | |
| if "api_key" not in st.session_state: | |
| api_key = st.text_input("Enter your API Key", type="password") | |
| if api_key: | |
| if api_key.isascii(): | |
| st.session_state.api_key = api_key | |
| client = openai.OpenAI(base_url = "https://integrate.api.nvidia.com/v1", | |
| api_key=api_key) | |
| st.rerun() | |
| else: | |
| st.warning("Please enter your API key correctly.") | |
| st.stop() | |
| else: | |
| st.warning("Please enter your API key to use the app. You can obtain your API key from here: https://build.nvidia.com/nvidia/nemotron-4-340b-instruct") | |
| st.stop() | |
| else: | |
| client = openai.OpenAI(base_url = "https://integrate.api.nvidia.com/v1", | |
| api_key=st.session_state.api_key) | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| encoding = tiktoken.get_encoding("cl100k_base") | |
| def num_tokens_from_string(string): | |
| num_tokens = len(encoding.encode(string)) | |
| return num_tokens | |
| def get_ai_response(prompt, chat_history): | |
| st.session_state.is_streaming = True | |
| st.session_state.response = "" | |
| system_tokens = num_tokens_from_string(system_prompt) | |
| user_tokens = num_tokens_from_string(prompt) | |
| # Maximum allowed tokens | |
| max_context_length = context_length - max_tokens | |
| available_tokens = max_context_length - system_tokens - user_tokens | |
| # Trim chat history if necessary | |
| trimmed_history = [] | |
| total_tokens = 0 | |
| for message in reversed(chat_history): | |
| message_tokens = num_tokens_from_string(message["content"]) | |
| if total_tokens + message_tokens <= available_tokens: | |
| trimmed_history.insert(0, message) | |
| total_tokens += message_tokens | |
| else: | |
| break | |
| # Construct final message list | |
| final_messages = [{"role": "system", "content": system_prompt}] + trimmed_history + [{"role": "user", "content": prompt}] | |
| try: | |
| with st.chat_message("assistant", avatar=st.session_state.assistant_avatar): | |
| stream = client.chat.completions.create( | |
| model=model, | |
| messages=final_messages, | |
| temperature=temperature, | |
| top_p=top_p, | |
| max_tokens=max_tokens, | |
| stream=True, | |
| ) | |
| placeholder = st.empty() | |
| with stylable_container( | |
| key="stop_generating", | |
| css_styles=""" | |
| button { | |
| position: fixed; | |
| bottom: 100px; | |
| left: 50%; | |
| transform: translateX(-50%); | |
| z-index: 1; | |
| } | |
| """, | |
| ): | |
| st.button("Stop generating") | |
| shown_message = "" | |
| for chunk in stream: | |
| if chunk.choices[0].delta.content is not None: | |
| content = chunk.choices[0].delta.content | |
| st.session_state.response += content | |
| shown_message += content.replace("\n", " \n")\ | |
| .replace("<", "\\<")\ | |
| .replace(">", "\\>") | |
| placeholder.markdown(shown_message) | |
| except: | |
| st.session_state.response += "\n!Exception\n" + traceback.format_exc() | |
| st.session_state.is_streaming = False | |
| return st.session_state.response | |
| def normalize_code_block(match): | |
| return match.group(0).replace(" \n", "\n")\ | |
| .replace("\\<", "<")\ | |
| .replace("\\>", ">") | |
| def normalize_inline(match): | |
| return match.group(0).replace("\\<", "<")\ | |
| .replace("\\>", ">") | |
| code_block_pattern = r"(```.*?```)" | |
| inline_pattern = r"`([^`\n]+?)`" | |
| def display_messages(): | |
| for i, message in enumerate(st.session_state.messages): | |
| 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: | |
| # Replace " \n" with "\n" within code blocks | |
| shown_message = re.sub(code_block_pattern, normalize_code_block, shown_message, flags=re.DOTALL) | |
| if "`" in shown_message: | |
| shown_message = re.sub(inline_pattern, normalize_inline, shown_message) | |
| st.markdown(shown_message) | |
| col1, col2, col3, col4 = st.columns([1, 1, 1, 1]) | |
| with col1: | |
| if st.button("Edit", key=f"edit_{i}_{len(st.session_state.messages)}"): | |
| st.session_state.edit_index = i | |
| st.rerun() | |
| with col2: | |
| if st.session_state.is_delete_mode and st.button("Delete", key=f"delete_{i}_{len(st.session_state.messages)}"): | |
| del st.session_state.messages[i] | |
| st.rerun() | |
| with col3: | |
| text_to_copy = message["content"] | |
| # Encode the string to escape | |
| text_to_copy_escaped = urllib.parse.quote(text_to_copy) | |
| copy_button_html = f""" | |
| <button id="copy-msg-btn-{i}" style='font-size: 1em; padding: 0.5em;' onclick='copyMessage("{i}")'>Copy</button> | |
| <script> | |
| function copyMessage(index) {{ | |
| navigator.clipboard.writeText(decodeURIComponent("{text_to_copy_escaped}")); | |
| let copyBtn = document.getElementById("copy-msg-btn-" + index); | |
| copyBtn.innerHTML = "Copied!"; | |
| setTimeout(function(){{ copyBtn.innerHTML = "Copy"; }}, 2000); | |
| }} | |
| </script> | |
| """ | |
| html(copy_button_html, height=50) | |
| if i == len(st.session_state.messages) - 1 and message["role"] == "CHATBOT": | |
| with col4: | |
| if st.button("Retry", key=f"retry_{i}_{len(st.session_state.messages)}"): | |
| if len(st.session_state.messages) >= 2: | |
| del st.session_state.messages[-1] | |
| st.session_state.retry_flag = True | |
| st.rerun() | |
| 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]["content"]) | |
| col1, col2 = st.columns([1, 1]) | |
| with col1: | |
| if st.form_submit_button("Save"): | |
| st.session_state.messages[i]["content"] = new_content | |
| del st.session_state.edit_index | |
| st.rerun() | |
| with col2: | |
| if st.form_submit_button("Cancel"): | |
| del st.session_state.edit_index | |
| st.rerun() | |
| # Add sidebar for advanced settings | |
| with st.sidebar: | |
| settings_tab, appearance_tab = st.tabs(["Settings", "Appearance"]) | |
| with settings_tab: | |
| st.markdown("Help (Japanese): https://rentry.org/9hgneofz") | |
| # Copy Conversation History button | |
| log_text = "" | |
| for message in st.session_state.messages: | |
| if message["role"] == "user": | |
| log_text += "<USER>\n" | |
| log_text += message["content"] + "\n\n" | |
| else: | |
| log_text += "<ASSISTANT>\n" | |
| log_text += message["content"] + "\n\n" | |
| log_text = log_text.rstrip("\n") | |
| # Encode the string to escape | |
| log_text_escaped = urllib.parse.quote(log_text) | |
| copy_log_button_html = f""" | |
| <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)); | |
| 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) | |
| if st.session_state.get("is_history_shown") != True: | |
| if st.button("Display History as Code Block"): | |
| st.session_state.is_history_shown = True | |
| st.rerun() | |
| else: | |
| if st.button("Hide History"): | |
| st.session_state.is_history_shown = False | |
| st.rerun() | |
| st.code(log_text) | |
| st.session_state.is_delete_mode = st.toggle("Enable Delete button") | |
| st.header("Advanced Settings") | |
| model = st.selectbox("Model", options=["nvidia/nemotron-4-340b-instruct"], index=0) | |
| system_prompt = st.text_area("System Prompt", height=200) | |
| temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.1) | |
| top_p = st.slider("Top-P", min_value=0.01, max_value=1.00, value=1.00, step=0.01) | |
| context_length = st.slider("Context Length", min_value=1, max_value=4096, value=4096, step=1) | |
| max_tokens = st.slider("Max Tokens (Output)", min_value=1, max_value=1024, value=1024, step=1) | |
| st.header("Restore History") | |
| history_input = st.text_area("Paste conversation history:", height=200) | |
| if st.button("Restore History"): | |
| st.session_state.messages = [] | |
| messages = re.split(r"^(<USER>|<ASSISTANT>)\n", history_input, flags=re.MULTILINE) | |
| role = None | |
| text = "" | |
| for message in messages: | |
| if message.strip() in ["<USER>", "<ASSISTANT>"]: | |
| if role and text: | |
| st.session_state.messages.append({"role": role, "content": text.strip()}) | |
| text = "" | |
| role = "user" if message.strip() == "<USER>" else "assistant" | |
| else: | |
| text += message | |
| if role and text: | |
| st.session_state.messages.append({"role": role, "content": text.strip()}) | |
| st.rerun() | |
| st.header("Clear History") | |
| if st.button("Clear Chat History"): | |
| st.session_state.messages = [] | |
| st.rerun() | |
| st.header("Change API Key") | |
| new_api_key = st.text_input("Enter new API Key", type="password") | |
| if st.button("Update API Key"): | |
| if new_api_key and new_api_key.isascii(): | |
| st.session_state.api_key = new_api_key | |
| client = openai.OpenAI(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", | |
| "Sawarabi Mincho": "Sawarabi Mincho" | |
| } | |
| selected_font = st.selectbox("Choose a font", ["Default"] + list(font_options.keys())) | |
| st.header("Change the font size") | |
| st.session_state.font_size = st.slider("Font size", min_value=16.0, max_value=50.0, value=16.0, step=1.0) | |
| st.header("Change the user's icon") | |
| st.session_state.user_avatar = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg", "webp", "gif", "bmp", "svg",], key="user_avatar_uploader") | |
| st.header("Change the assistant's icon") | |
| st.session_state.assistant_avatar = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg", "webp", "gif", "bmp", "svg",], key="assistant_avatar_uploader") | |
| st.header("Change the icon size") | |
| st.session_state.avatar_size = st.slider("Icon size", min_value=2.0, max_value=20.0, value=2.0, step=0.2) | |
| # After Stop generating | |
| if st.session_state.get("is_streaming"): | |
| 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 | |
| st.rerun() | |
| if selected_font != "Default": | |
| with open("style.css") as css: | |
| st.markdown(f'<style>{css.read()}</style>', unsafe_allow_html=True) | |
| st.markdown(f'<style>body * {{ font-family: "{font_options[selected_font]}", serif !important; }}</style>', unsafe_allow_html=True) | |
| # Change font size | |
| st.markdown(f'<style>[data-testid="stChatMessageContent"] .st-emotion-cache-cnbvxy p{{font-size: {st.session_state.font_size}px;}}</style>', unsafe_allow_html=True) | |
| # Change icon size | |
| # (CSS element names may be subject to change.) | |
| # (Contributor: ★31 >>538) | |
| AVATAR_SIZE_STYLE = f""" | |
| <style> | |
| [data-testid="chatAvatarIcon-user"] {{ | |
| width: {st.session_state.avatar_size}rem; | |
| height: {st.session_state.avatar_size}rem; | |
| }} | |
| [data-testid="chatAvatarIcon-assistant"] {{ | |
| width: {st.session_state.avatar_size}rem; | |
| height: {st.session_state.avatar_size}rem; | |
| }} | |
| [data-testid="stChatMessage"] .st-emotion-cache-1pbsqtx {{ | |
| width: {st.session_state.avatar_size / 1.6}rem; | |
| height: {st.session_state.avatar_size / 1.6}rem; | |
| }} | |
| [data-testid="stChatMessage"] .st-emotion-cache-p4micv {{ | |
| width: {st.session_state.avatar_size}rem; | |
| height: {st.session_state.avatar_size}rem; | |
| }} | |
| </style> | |
| """ | |
| st.markdown(AVATAR_SIZE_STYLE, unsafe_allow_html=True) | |
| display_messages() | |
| # After Retry | |
| if st.session_state.get("retry_flag"): | |
| if len(st.session_state.messages) > 0: | |
| prompt = st.session_state.messages[-1]["content"] | |
| messages = st.session_state.messages[:-1].copy() | |
| response = get_ai_response(prompt, 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..."): | |
| chat_history = st.session_state.messages.copy() | |
| shown_message = prompt.replace("\n", " \n")\ | |
| .replace("<", "\\<")\ | |
| .replace(">", "\\>") | |
| with st.chat_message("user", avatar=st.session_state.user_avatar): | |
| st.write(shown_message) | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| response = get_ai_response(prompt, chat_history) | |
| st.session_state.messages.append({"role": "assistant", "content": response}) | |
| st.rerun() |