Spaces:
Sleeping
Sleeping
File size: 5,009 Bytes
86ce226 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
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"])
img = Image.open(io.BytesIO(img_data))
return img, "OK"
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:
gr.Markdown("# 🔬 LLM-Microscope — Understanding Token Representations in Transformers")
gr.Markdown("Select a model, a mode of analysis, and a sentence. The tool will visualize what’s happening **inside** the language model — layer by layer, token by token.")
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 input text:", value="I love to live my life")
submit = gr.Button("Submit")
box_for_plot = gr.Image(label="Visualization", type="pil")
# 💬 Explanation below the visualization
explanation_text = gr.Markdown("""
### 📘 Legend and Interpretation
This heatmap shows **how each token is processed** across layers of a language model. Here's how to read it:
- **Rows**: layers of the model (bottom = deeper)
- **Columns**: input tokens
- **Colors**: intensity of effect (depends on the selected metric)
**Metrics explained:**
- `Layer wise non-linearity`: how nonlinear the transformation is at each layer (red = more nonlinear).
- `Next-token prediction from intermediate representations`: shows which layers begin to make good predictions.
- `Contextualization measurement`: tokens with more context info get lower scores (green = more context).
- `Layerwise predictions (logit lens)`: tracks how the model’s guesses evolve at each layer.
- `Tokenwise loss without i-th layer`: shows how much each token depends on a specific layer. Red means performance drops if we skip this layer.
Use this tool to **peek inside the black box** — it reveals which layers matter most, which tokens carry the most memory, and how LLMs evolve their predictions.
""")
def update_output(task_name: str, model_name: str, text: str, normalization_type: str) -> Tuple[Any]:
img, _ = get_layerwise_nonlinearity(task_name, model_name, text, normalization_type)
return img
def set_default(task_name: str) -> str:
if task_name in ["Layer wise non-linearity", "Next-token prediction from intermediate representations", "Tokenwise loss without i-th layer"]:
return "token-wise"
return "global"
def check_normalization(task_name: str, normalization_name) -> Tuple[str]:
if task_name == "Contextualization measurement" and normalization_name == "token-wise":
return "global"
return normalization_name
task_selector.select(set_default, [task_selector], [normalization_selector])
normalization_selector.select(check_normalization, [task_selector, normalization_selector], [normalization_selector])
submit.click(
fn=update_output,
inputs=[task_selector, model_selector, text_message, normalization_selector],
outputs=[box_for_plot]
)
if __name__ == "__main__":
demo.launch(share=True, server_port=7860, server_name="0.0.0.0")
|