Update app.py

app.py

@@ -2,35 +2,85 @@ import gradio as gr
 import requests
 from PIL import Image
 import io
-from typing import Any
+from typing import Any, Tuple
 import os
 
 class Client:
     def __init__(self, server_url: str):
         self.server_url = server_url
 
-    def send_request(self, model_name: str, text: str) -> Any:
-        response = requests.post(self.server_url, json={"model_name": model_name, "text": text})
+    def send_request(self, task_name: str, model_name: str, text: str, normalization_type: str) -> Tuple[Any, str]:
+        response = requests.post(self.server_url, json={"task_name": task_name, "model_name": model_name, "text": text, "normalization_type": normalization_type}, timeout=10)
         if response.status_code == 200:
-            img_data = response.content
+            response_data = response.json()
+            img_data = bytes.fromhex(response_data["image"])
+            log_info = response_data["log"]
             img = Image.open(io.BytesIO(img_data))
-            return img
+            return img, log_info
         else:
-            return "Error, please retry"
+            return "Error, please retry", "Error: Could not get response from server"
 
 client = Client(f"http://{os.environ['SERVER']}/predict")
 
-def get_layerwise_nonlinearity(model_name: str, text: str) -> Any:
-    return client.send_request(model_name, text)
+def get_layerwise_nonlinearity(task_name: str, model_name: str, text: str, normalization_type: str) -> Tuple[Any, str]:
+    return client.send_request(task_name, model_name, text, normalization_type)
 
 with gr.Blocks() as demo:
+    with gr.Row():
+        model_selector = gr.Dropdown(
+            choices=[
+                "facebook/opt-1.3b",
+                "facebook/opt-2.7b",
+                "microsoft/Phi-3-mini-128k-instruct"
+            ],
+            value="facebook/opt-1.3b",
+            label="Select Model"
+        )
+        task_selector = gr.Dropdown(
+            choices=[
+                "Layer wise non-linearity (with first layer)", 
+                "Next-token prediction from intermediate representations", 
+                "Contextualization mesurment",
+                "Layerwise predictions and losses",
+                "Tokenwise loss without i-th layer"
+            ], 
+            value="Layer wise non-linearity (with first layer)",
+            label="Select Mode"
+        )
+        normalization_selector = gr.Dropdown(
+            choices=["global", "token-wise"], #, "sentence-wise"],
+            value="token-wise",
+            label="Select Normalization"
+        )
     with gr.Column():
-        model_selector = gr.Dropdown(choices=["mistralai/Mistral-7B-v0.1", "facebook/opt-125m"], label="Select Model")
-        text_message = gr.Textbox(label="Enter your request:")
+        text_message = gr.Textbox(label="Enter your request:", value="I love to live my life")
         submit = gr.Button("Submit")
-        box_for_plot = gr.Image(label="Layer wise non-linearity (with first layer)", type="pil")
+        box_for_plot = gr.Image(label="Visualization", type="pil")
+        log_output = gr.Textbox(label="Log Output", lines=10, interactive=False, value="")
 
-        submit.click(get_layerwise_nonlinearity, [model_selector, text_message], [box_for_plot])
+        def update_output(task_name: str, model_name: str, text: str, normalization_type: str, existing_log: str) -> Tuple[Any, str]:
+            img, new_log = get_layerwise_nonlinearity(task_name, model_name, text, normalization_type)
+            combined_log = existing_log + "---\n" + new_log + "\n"
+            return img, combined_log
+        
+        def set_default(task_name: str) -> str:
+            if task_name == "Layer wise non-linearity (with first layer)":
+                return "token-wise"
+            if task_name == "Next-token prediction from intermediate representations":
+                return "token-wise"
+            if task_name == "Contextualization mesurment":
+                return "global"
+            if task_name == "Layerwise predictions and losses":
+                return "global"
+            if task_name == "Tokenwise loss without i-th layer":
+                return "token-wise"
+
+        task_selector.select(set_default, [task_selector], [normalization_selector])
+        submit.click(
+            fn=update_output, 
+            inputs=[task_selector, model_selector, text_message, normalization_selector, log_output], 
+            outputs=[box_for_plot, log_output]
+        )
 
 if __name__ == "__main__":
     demo.launch(share=True, server_port=7860, server_name="0.0.0.0")