| import gradio as gr |
| import requests |
| from PIL import Image |
| import io |
| from typing import Any, Tuple |
| import os |
|
|
| class Client: |
| def __init__(self, server_url: str): |
| self.server_url = server_url |
|
|
| 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=60 |
| ) |
| if response.status_code == 200: |
| 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, log_info |
| else: |
| return "Error, please retry", "Error: Could not get response from server" |
|
|
| client = Client(f"http://{os.environ['SERVER']}/predict") |
|
|
| 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", |
| "TheBloke/Llama-2-7B-fp16" |
| |
| |
| ], |
| value="facebook/opt-1.3b", |
| label="Select Model" |
| ) |
| task_selector = gr.Dropdown( |
| choices=[ |
| "Layer wise non-linearity", |
| "Next-token prediction from intermediate representations", |
| "Contextualization measurement", |
| "Layerwise predictions (logit lens)", |
| "Tokenwise loss without i-th layer" |
| ], |
| value="Layer wise non-linearity", |
| label="Select Mode" |
| ) |
| normalization_selector = gr.Dropdown( |
| choices=["global", "token-wise"], |
| value="token-wise", |
| label="Select Normalization" |
| ) |
| with gr.Column(): |
| 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="Visualization", type="pil") |
| log_output = gr.Textbox(label="Log Output", lines=10, interactive=False, value="") |
|
|
| 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": |
| return "token-wise" |
| if task_name == "Next-token prediction from intermediate representations": |
| return "token-wise" |
| if task_name == "Contextualization measurement": |
| return "global" |
| if task_name == "Layerwise predictions (logit lens)": |
| return "global" |
| if task_name == "Tokenwise loss without i-th layer": |
| return "token-wise" |
|
|
| def check_normalization(task_name: str, normalization_name) -> Tuple[str, str]: |
| if task_name == "Contextualization measurement" and normalization_name == "token-wise": |
| return ("global", "\nALERT: Cannot apply token-wise normalization to one sentence, setting global normalization\n") |
| return (normalization_name, "") |
| |
| task_selector.select(set_default, [task_selector], [normalization_selector]) |
| normalization_selector.select(check_normalization, [task_selector, normalization_selector], [normalization_selector, log_output]) |
| 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") |
|
|
