Basic app using gpt2 working

app.py

@@ -1,17 +1,89 @@
-from platform import python_version
+import torch
+from tuned_lens.nn.lenses import TunedLens, LogitLens
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from tuned_lens.nn import TunedLens
 from tuned_lens.plotting import plot_lens
-
 import gradio as gr
+from plotly import graph_objects as go
+
+device = torch.device("cpu")
+print(f"Using device {device} for inference")
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+model = model.to(device)
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tuned_lens = TunedLens.load("lens/gpt2", map_location=device)
+logit_lens = LogitLens(model)
+
+lens_options_dict = {
+    "Tuned Lens": tuned_lens,
+    "Logit Lens": logit_lens,
+}
+statistic_options_dict = {
+    "Entropy": "entropy",
+    "Cross Entropy": "ce",
+    "Forward KL": "forward_kl",
+}
+
+
+def make_plot(lens, text, statistic, token_cutoff):
+    input_ids = tokenizer.encode(text, return_tensors="pt")
+
+    if len(input_ids[0]) == 0:
+        return go.Figure(layout=dict(title="Please enter some text."))
+
+    if token_cutoff < 1:
+        return go.Figure(layout=dict(title="Please provide valid token cut off."))
+
+    fig = plot_lens(
+        model,
+        tokenizer,
+        lens_options_dict[lens],
+        layer_stride=1,
+        input_ids=input_ids,
+        start_pos=max(len(input_ids[0]) - token_cutoff, 0),
+        statistic=statistic_options_dict[statistic],
+    )
+    fig.update_layout(template="plotly_dark")
+
+    # Update the colorscale of the heatmap trace
+    for trace in fig.data:
+        if trace.type == "heatmap":
+            trace.update(colorscale="Inferno")
+
+    return fig
+
+
+preamble = """
+# The Tuned Lens 🔎
+
+A tuned lens allows us to peak at the iterative computations a transformer uses to compute the next token.
+
+A lens into a transformer with n layers allows you to replace the last $m$ layers of the model with an [affine transformation](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) (we call these affine translators).
+
+This essentially skips over these last few layers and lets you see the best prediction that can be made from the model's representations, i.e. the residual stream, at layer $n - m$. Since the representations may be rotated, shifted, or stretched from layer to layer it's useful to train the len's affine adapters specifically on each layer. This training is what differentiates this method from simpler approaches that decode the residual stream of the network directly using the unembeding layer i.e. the logit lens. We explain this process in [the paper](https://arxiv.org/abs/2303.08112).
+"""
 
-LENS_PATH = '<PATH TO LENS>'
 
-def plot_lens_outputs(text):
-    model = AutoModelForCausalLM.from_pretrained('gpt2')
-    tokenizer = AutoTokenizer.from_pretrained('gpt2')
-    #lens = TunedLens.load(LENS_PATH)
-    return gr.outputs.Plot(plot_lens(model, tokenizer, text=text))
+with gr.Blocks() as iface:
+    gr.Markdown(preamble)
+    with gr.Column():
+        text = gr.Textbox(
+            value="the iterative computations a transformer uses to compute the next",
+            label="Input Text",
+        )
+        with gr.Row():
+            lens_options = gr.Dropdown(
+                list(lens_options_dict.keys()), value="Tuned Lens", label="Select Lens"
+            )
+            statistic = gr.Dropdown(
+                list(statistic_options_dict.keys()),
+                value="Entropy",
+                label="Select Statistic",
+            )
+            token_cutoff = gr.Slider(
+                maximum=20, minimum=2, value=10, label="Token Cut Off"
+            )
+        examine_btn = gr.Button(value="Examine")
+        plot = gr.Plot()
+    examine_btn.click(make_plot, [lens_options, text, statistic, token_cutoff], plot)
 
-iface = gr.Interface(fn=plot_lens_outputs, inputs="text", outputs=gr.outputs.Plot(type="auto"))
 iface.launch()