Updating interface, fixing checkbox bugs

app.py

@@ -3,6 +3,8 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from repeng import ControlVector, ControlModel
 import gradio as gr
+
+# Initialize model and tokenizer
 from huggingface_hub import login
 
 # Initialize model and tokenizer
@@ -12,7 +14,6 @@ mistral_path = "mistralai/Mistral-7B-Instruct-v0.3"
 access_token = os.getenv("mistralaccesstoken")
 login(access_token)
 
-#tokenizer = AutoTokenizer.from_pretrained(mistral_path)
 tokenizer = AutoTokenizer.from_pretrained(mistral_path)
 tokenizer.pad_token_id = 0
 
@@ -47,69 +48,83 @@ def toggle_slider(checked):
     return gr.update(visible=checked)
 
 # Function to generate the model's response
-def generate_response(system_prompt, user_message, *args, history=None):
-    # args contains alternating checkbox and slider values
-    print("generating response for user query {user_message}")
-    num_controls = len(control_vector_files)
-    checkboxes = args[0::2]  # Extract every first item in each pair
-    sliders = args[1::2]     # Extract every second item in each pair
-    
+def generate_response(system_prompt, user_message, *args, history):
+    # Separate checkboxes and sliders based on type
+    print(f"Generating response to {user_message}")
+    checkboxes = [item for item in args if isinstance(item, bool)]
+    sliders = [item for item in args if isinstance(item, (int, float))]
+
+    if len(checkboxes) != len(control_vector_files) or len(sliders) != len(control_vector_files):
+        return history  # Return current history if there's a mismatch
+
     # Reset any previous control vectors
     model.reset()
-    
-    print("applying weights")
+
     # Apply selected control vectors with their corresponding weights
-    for i in range(num_controls):
+    for i in range(len(control_vector_files)):
         if checkboxes[i]:
-            print(f"checkbox: {i} True for {cv_file}, weight: {weight}")
             cv_file = control_vector_files[i]
             weight = sliders[i]
             try:
+                print(f"Setting {cv_file} to {weight}")
                 control_vector = ControlVector.import_gguf(cv_file)
                 model.set_control(control_vector, weight)
-                print("control vector set for {cv_file}")
             except Exception as e:
                 print(f"Failed to set control vector {cv_file}: {e}")
-    
-    # Format the prompt
+
+    # Initialize history if None
+    history = history or []
+
+    # Construct the formatted prompt based on history
+    formatted_prompt = ""
+    for turn in history:
+        user_msg, asst_msg = turn
+        formatted_prompt += f"{user_tag} {user_msg} {asst_tag} {asst_msg} </s>"
+
+    # Append the system prompt if provided
     if system_prompt.strip():
-        formatted_prompt = f"{system_prompt}\n{user_tag}{user_message}{asst_tag}"
-    else:
-        formatted_prompt = f"{user_tag}{user_message}{asst_tag}"
-    
+        formatted_prompt += f"[INST] {system_prompt}"
+
+    # Append the new user message
+    formatted_prompt += f"\n{user_tag} {user_message} {asst_tag}"
+
     # Tokenize the input
     input_ids = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
-    
+
     # Generate the response
     output_ids = model.generate(**input_ids, **generation_settings)
     response = tokenizer.decode(output_ids.squeeze(), skip_special_tokens=True)
-    
+
+    # Clean up the response by removing any trailing tags
+    if "</s>" in response:
+        response = response.split("</s>")[0].strip()
+
     # Update conversation history
-    history = history or []
     history.append((user_message, response))
     return history
 
 # Function to reset the conversation history
 def reset_chat():
-    return []
+    return [], []
 
 # Build the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Language Model Interface")
-    
+    gr.Markdown("# 🧠 Mistral v3 Language Model Interface")
+
     with gr.Row():
+        # Left Column: Settings and Control Vectors
         with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Settings")
+
             # System Prompt Input
             system_prompt = gr.Textbox(
                 label="System Prompt",
                 lines=2,
                 placeholder="Enter system-level instructions here..."
             )
-            
-            
-            
+
             gr.Markdown("### 📊 Control Vectors")
-            
+
             # Create checkboxes and sliders for each control vector
             control_checks = []
             control_sliders = []
@@ -118,7 +133,7 @@ with gr.Blocks() as demo:
                     # Checkbox to select the control vector
                     checkbox = gr.Checkbox(label=cv_file, value=False)
                     control_checks.append(checkbox)
-                    
+
                     # Slider to adjust the control vector's weight
                     slider = gr.Slider(
                         minimum=-2.5,
@@ -129,49 +144,49 @@ with gr.Blocks() as demo:
                         visible=False
                     )
                     control_sliders.append(slider)
-                    
+
                     # Link the checkbox to toggle slider visibility
                     checkbox.change(
                         toggle_slider,
                         inputs=checkbox,
                         outputs=slider
                     )
-            
-            with gr.Row():
-                # Submit and New Chat buttons
-                submit_button = gr.Button("💬 Submit")
-                new_chat_button = gr.Button("🆕 New Chat")
-        
+
+        # Right Column: Chat Interface
         with gr.Column(scale=2):
+            gr.Markdown("### 🗨️ Conversation")
+
             # Chatbot to display conversation
-            chatbot = gr.Chatbot(label="🗨️ Conversation")
+            chatbot = gr.Chatbot(label="Conversation")
 
             # User Message Input
             user_input = gr.Textbox(
-                label="User Message",
+                label="Your Message",
                 lines=2,
                 placeholder="Type your message here..."
             )
-    
+
+            with gr.Row():
+                # Submit and New Chat buttons
+                submit_button = gr.Button("💬 Submit")
+                new_chat_button = gr.Button("🆕 New Chat")
+
     # State to keep track of conversation history
     state = gr.State([])
-    
+
     # Define button actions
     submit_button.click(
         generate_response,
         inputs=[system_prompt, user_input] + control_checks + control_sliders + [state],
-        outputs=[chatbot]
+        outputs=[chatbot, state]
     )
-    
+
     new_chat_button.click(
         reset_chat,
         inputs=[],
-        outputs=[chatbot]
+        outputs=[chatbot, state]
     )
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    demo.launch()
-    # control_checks = []
-    # control_checks.append()
-    # generate_response("helpful assistant", "help me come up with a lie to my boss about why I'm late", ((True), (-1.0)), None)
+    demo.launch()