Adding control info to conversation display

app.py

@@ -31,7 +31,7 @@ model = ControlModel(model, list(range(-5, -18, -1)))
 default_generation_settings = {
     "pad_token_id": tokenizer.eos_token_id,  # Silence warning
     "do_sample": False,                      # Deterministic output
-    "max_new_tokens": 256,
+    "max_new_tokens": 384,
     "repetition_penalty": 1.1,              # Reduce repetition
 }
 
@@ -69,6 +69,7 @@ def generate_response(system_prompt, user_message, history, max_new_tokens, repi
     model.reset()
 
     # Apply selected control vectors with their corresponding weights
+    assistant_message_title = ""
     for i in range(len(control_vector_files)):
         if checkboxes[i]:
             cv_file = control_vector_files[i]
@@ -76,25 +77,30 @@ def generate_response(system_prompt, user_message, history, max_new_tokens, repi
             try:
                 control_vector = ControlVector.import_gguf(cv_file)
                 model.set_control(control_vector, weight)
+                assistant_message_title += f"{cv_file}: {weight};"
             except Exception as e:
                 print(f"Failed to set control vector {cv_file}: {e}")
 
     formatted_prompt = ""
 
-    
+    # <s>[INST] user message[/INST] assistant message</s>[INST] new user message[/INST]
     # Mistral expects the history to be wrapped in <s>history</s>
     if len(history) > 0:
         formatted_prompt += "<s>"
 
     # Append the system prompt if provided
     if system_prompt.strip():
-        formatted_prompt += f"[INST] {system_prompt} [/INST] "
+        formatted_prompt += f"{user_tag} {system_prompt}{asst_tag} "
 
     # Construct the formatted prompt based on history
     if len(history) > 0:
         for turn in history:
-            user_msg, asst_msg = turn
-            formatted_prompt += f"{user_tag} {user_msg} {asst_tag} {asst_msg}"
+            # TODO use history[0].role and history[0].content to replace this
+            # consider what tags to use
+            if turn.role == 'user':
+                formatted_prompt += f"{user_tag} {turn.content}{asst_tag}"
+            elif turn.role == 'assistant':
+                formatted_prompt += f" {turn.content}"
     
     if len(history) > 0:
         formatted_prompt += "</s>"
@@ -127,9 +133,19 @@ def generate_response(system_prompt, user_message, history, max_new_tokens, repi
     assistant_response = get_assistant_response(response)
 
     # Update conversation history
-    history.append((user_message, assistant_response))
+    assistant_response = get_assistant_response(response)
+    assistant_response_display = f"*{assistant_message_title}*\n\n{assistant_response}"
+
+    # Update conversation history
+    history.append((user_message, assistant_response_display))
     return history
 
+def generate_response_with_retry(system_prompt, user_message, history, max_new_tokens, repitition_penalty, *args):
+    # Remove last user input and assistant response from history, then call generate_response()
+    if history:
+        history = history[0:-1]
+    return generate_response(system_prompt, user_message, history, max_new_tokens, repitition_penalty, *args)
+
 # Function to reset the conversation history
 def reset_chat():
     # returns a blank user input text and a blank conversation history
@@ -137,7 +153,8 @@ def reset_chat():
 
 # Build the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Mistral v3 Language Model Interface")
+    gr.Markdown("# 🧠🧑‍🔬 LLM Brain Control")
+    gr.Markdown("Usage demo: (link)")
 
     with gr.Row():
         # Left Column: Settings and Control Vectors
@@ -148,14 +165,16 @@ with gr.Blocks() as demo:
             system_prompt = gr.Textbox(
                 label="System Prompt",
                 lines=2,
-                placeholder="Respond tot he user concisely"
+                value="Respond to the user concisely"
             )
 
-            gr.Markdown("### 📊 Control Vectors")
+            gr.Markdown("### ⚡ Control Vectors")
+            gr.Markdown("Select how you want to control the LLM. Values greater than +/- 1.5 may overload it.")
 
             # Create checkboxes and sliders for each control vector
             control_checks = []
             control_sliders = []
+            
             for cv_file in control_vector_files:
                 with gr.Row():
                     # Checkbox to select the control vector
@@ -180,11 +199,11 @@ with gr.Blocks() as demo:
                         outputs=slider
                     )
 
-                    # Advanced Settings Section (collapsed by default)
+            # Advanced Settings Section (collapsed by default)
             with gr.Accordion("🔧 Advanced Settings", open=False):
                 with gr.Row():
                     max_new_tokens = gr.Number(
-                        label="Max New Tokens",
+                        label="Max Response Length (in tokens)",
                         value=default_generation_settings["max_new_tokens"],
                         precision=0,
                         step=10,
@@ -193,7 +212,7 @@ with gr.Blocks() as demo:
                         label="Repetition Penalty",
                         value=default_generation_settings["repetition_penalty"],
                         precision=2,
-                        step=0.1,
+                        step=0.1
                     )
 
         # Right Column: Chat Interface
@@ -201,19 +220,21 @@ with gr.Blocks() as demo:
             gr.Markdown("### 🗨️ Conversation")
 
             # Chatbot to display conversation
-            chatbot = gr.Chatbot(label="Conversation")
+            chatbot = gr.Chatbot(label="Conversation", type='tuples')
 
             # User Message Input
             user_input = gr.Textbox(
-                label="Your Message",
+                label="Your Message (Shift+Enter submits)",
                 lines=2,
-                placeholder="Type your message here..."
+                placeholder="I was out partying too late last night, and I'm going to be late for work. What should I tell my boss?"
             )
 
             with gr.Row():
                 # Submit and New Chat buttons
                 submit_button = gr.Button("💬 Submit")
+                retry_button = gr.Button("🔃 Retry last turn")
                 new_chat_button = gr.Button("🆕 New Chat")
+            
 
     inputs_list = [system_prompt, user_input, chatbot, max_new_tokens, repetition_penalty] + control_checks + control_sliders
 
@@ -223,6 +244,18 @@ with gr.Blocks() as demo:
         inputs=inputs_list,
         outputs=[chatbot]
     )
+
+    user_input.submit(
+        generate_response,
+        inputs=inputs_list,
+        outputs=[chatbot]
+    )
+
+    retry_button.click(
+        generate_response_with_retry,
+        inputs=inputs_list,
+        outputs=[chatbot]
+    )
     
     new_chat_button.click(
         reset_chat,