LLM-Microscope / app.py
matveymih's picture
Update app.py
5dfe7bb verified
raw
history blame
4.39 kB
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"
# "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",
"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"], #, "sentence-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")