Add the Prefill and the @prefill command

app.py

@@ -33,9 +33,18 @@ else:
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+if "prefill" not in st.session_state:
+    st.session_state.prefill = ""
+
 def get_ai_response(messages):
     st.session_state.is_streaming = True
     st.session_state.response = ""
+    shown_message = ""
+
+    if st.session_state.prefill:
+        messages.append({"role": "assistant", "content": st.session_state.prefill})
+        st.session_state.response += st.session_state.prefill
+        shown_message = st.session_state.prefill.replace("\n", "  \n")
 
     with st.chat_message("assistant"):
         penalty_kwargs = {
@@ -69,8 +78,6 @@ def get_ai_response(messages):
         ):
             st.button("Stop generating")
 
-        shown_message = ""
-
         for chunk in stream:
             if chunk.choices and chunk.choices[0].delta.get("content"):
                 content = chunk.choices[0].delta["content"]
@@ -78,6 +85,11 @@ def get_ai_response(messages):
                 shown_message += content.replace("\n", "  \n")
                 placeholder.markdown(shown_message)
 
+    if st.session_state.prefill == st.session_state.response:
+        st.session_state.is_error = True
+    else:
+        st.session_state.is_error = False
+
     st.session_state.is_streaming = False
     return st.session_state.response
 
@@ -205,8 +217,8 @@ with st.sidebar:
                   "llama-2-70b-chat",
                   "llama-2-7b-chat",
                   "llama-2-70b-fast",
-                  "pplx-70b-chat",
-                  "pplx-70b-online",
+                  "llama-3-70b-instruct",
+                  "llama-3-8b-instruct",
                   "mistral-7b-instruct-v0.1",
                   "mixtral-8x7b-instruct-v0.1",
                   "mixtral-8x7b-fast",
@@ -216,8 +228,6 @@ with st.sidebar:
                   "mistral-tiny",
                   "dolphin-mixtral-8x7b",
                   "mixtral-8x22b",
-                  "pplx-7b-chat",
-                  "pplx-7b-online",
                   "yi-34-chat",
                   "chronos-hermes-13b",
                   "mythomax-l2-13b",
@@ -229,6 +239,10 @@ with st.sidebar:
                   ]
     model = st.selectbox("Model", options=model_list, index=0)
     system_prompt = st.text_area("System prompt", height=200)
+
+    st.session_state.prefill = st.text_area("Prefill", height=50, value=st.session_state.prefill, placeholder="It only works well with the Claude models.")
+    save_prefill = st.toggle("Save the @prefill command input in the sidebar", value=True)
+    
     temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=1.0, step=0.1)
     top_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"])
@@ -290,14 +304,10 @@ display_messages()
 # After Retry
 if "retry_flag" in st.session_state and st.session_state.retry_flag == True:
     if len(st.session_state.messages) > 0:
-        response = get_ai_response(st.session_state.messages)
+        messages = st.session_state.messages.copy()
+        response = get_ai_response(messages)
         st.session_state.messages.append({"role": "assistant", "content": response})
 
-        if response:
-            st.session_state.is_error = False
-        else:
-            st.session_state.is_error = True
-
         st.session_state.retry_flag = False
         st.rerun()
             
@@ -305,6 +315,17 @@ if "retry_flag" in st.session_state and st.session_state.retry_flag == True:
         st.session_state.retry_flag = False
 
 if prompt := st.chat_input("What is up?"):
+    used_prefill = False
+    prefill_pattern = r"([@＠](prefill|ぷれふぃる|プレフィル)\s?(.*))"
+    prefill_match = re.search(prefill_pattern, prompt)
+
+    if prefill_match:
+        used_prefill = True
+        if not save_prefill:
+            original_prefill = st.session_state.prefill
+        st.session_state.prefill = prefill_match.group(3)
+        prompt = prompt.replace(prefill_match.group(1), '')
+
     st.session_state.messages.append({"role": "user", "content": prompt})
     messages = st.session_state.messages.copy()
 
@@ -313,11 +334,9 @@ if prompt := st.chat_input("What is up?"):
 
     response = get_ai_response(messages)
 
-    if response:
-        st.session_state.is_error = False
-    else:
-        st.session_state.is_error = True
-
     st.session_state.messages.append({"role": "assistant", "content": response})
+
+    if used_prefill and not save_prefill:
+        st.session_state.prefill = original_prefill
     
     st.rerun()