Update app.py

app.py

@@ -5,23 +5,20 @@ import pandas as pd
 import io
 
 # ----------------- Configure Your Server API Address -----------------
-# Replace the IP address and port with your own
 SERVER_URL = "http://103.235.229.133:7000/predict"
 # ---------------------------------------------------------------------
 
 # ----------------- Fixed Model Name -----------------
-# As per requirements, the model name is fixed and not a UI input.
 FIXED_MODEL_NAME = "Qwen2.5-1.5B-Instruct"
 # ----------------------------------------------------
 
 # Default input data for easy testing
 DEFAULT_DATAFRAME_VALUE = [
     ["France", "capital city of", "Paris", "she traveled to France for the first time"],
-    # ["India", "official currency of", "Rupee", "caves of southern India and similar evidence"],
-    # ["Australia", "largest city in", "Sydney", "Grainger left Australia at the age of 13"]
 ]
 
-def call_api(only_final_result, data_df):
+# --- KEY CHANGE 1: Update the function signature to accept the new parameter ---
+def call_api(patched_layer_number, only_final_result, data_df):
     """
     This function is called by Gradio. It collects inputs from the UI,
     formats them for the API, and sends the request.
@@ -38,19 +35,18 @@ def call_api(only_final_result, data_df):
     data_df_with_id.to_csv(output, index=False)
     csv_data = output.getvalue()
     
-    # Prepare the payload for the API
+    # --- KEY CHANGE 2: Add the new parameter to the API payload ---
     payload = {
         "data_csv": csv_data,
+        "patched_layer_number": patched_layer_number,
         "only_final_result": only_final_result
     }
 
     try:
-        # Send the request to the server
         response = requests.post(SERVER_URL, json=payload, timeout=300, stream=True)
-        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
+        response.raise_for_status()
         result_json = response.json()
         
-        # Extract the result data for the DataFrame
         dataframe_result_data = result_json.get("result", [])
         
         if isinstance(dataframe_result_data, list) and len(dataframe_result_data) > 0:
@@ -83,10 +79,19 @@ with gr.Blocks(theme=gr.themes.Glass(), css=".gradio-container {max-width: 1280p
         gr.Markdown("### 1. Configure Parameters")
         with gr.Group():
             gr.Markdown(f"**Current Model (Fixed):** `{FIXED_MODEL_NAME}`")
+            
+            # --- KEY CHANGE 3: Add the new Gradio input component ---
+            patched_layer_number_input = gr.Number(
+                label="Patched Layer Number (patched_layer_number)",
+                value=15,  # Set a reasonable default value
+                precision=0,  # This ensures the input is an integer
+                info="Specify the target layer number for patching."
+            )
+            
             only_final_result_input = gr.Checkbox(
                 label="Only Return Final Result (only_final_result)",
                 value=False,
-                info="If checked, the API will only return the final aggregated result, not the detailed information for all layers."
+                info="If checked, the API will only return the final aggregated result, not detailed info for all layers."
             )
         
         gr.Markdown("### 2. Input Data")
@@ -96,8 +101,6 @@ with gr.Blocks(theme=gr.themes.Glass(), css=".gradio-container {max-width: 1280p
             label="Data Table",
             headers=["subject", "relation", "object", "prompt_source"],
             value=DEFAULT_DATAFRAME_VALUE,
-            # Set the initial number of rows to 5 to give the table more vertical space on load.
-            # (rows_to_display, "dynamic" allows user to add/remove rows)
             row_count=(5, "dynamic"),
             col_count=(4, "fixed"),
             interactive=True
@@ -112,17 +115,20 @@ with gr.Blocks(theme=gr.themes.Glass(), css=".gradio-container {max-width: 1280p
                 output_dataframe = gr.DataFrame(
                     label="Inference Result Details", 
                     interactive=False,
+                    # You might want to update headers if the output changes based on this parameter
                     headers=["id", "subject", "relation", "object", "prompt_source", "prompt_target", "layer_source", "is_correct_patched", "generations"]
                 )
             with gr.TabItem("Raw JSON Response"):
                 output_json = gr.JSON(label="Full JSON response from API")
             
+    # --- KEY CHANGE 4: Update the click event to include the new input ---
+    # The order of inputs here MUST match the order of arguments in the `call_api` function.
     submit_btn.click(
         fn=call_api,
-        inputs=[only_final_result_input, data_input],
+        inputs=[patched_layer_number_input, only_final_result_input, data_input],
         outputs=[output_json, output_dataframe]
     )
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()