upgraded app to use tuned_lens=0.1.0

README.md

@@ -3,6 +3,7 @@ title: Tuned Lens
 emoji: 🔎
 colorFrom: pink
 colorTo: blue
+port: 7860
 sdk: docker
 pinned: false
 license: mit

app.py

@@ -1,17 +1,20 @@
 import torch
 from tuned_lens.nn.lenses import TunedLens, LogitLens
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from tuned_lens.plotting import plot_lens
+from tuned_lens.plotting import PredictionTrajectory
 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("EleutherAI/pythia-410m-deduped-v0")
+model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-410m-deduped")
 model = model.to(device)
-tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-410m-deduped-v0")
-tuned_lens = TunedLens.load("pythia-410m-deduped-v0", map_location=device)
-logit_lens = LogitLens(model)
+tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-410m-deduped")
+tuned_lens = TunedLens.from_model_and_pretrained(
+    model=model,
+    map_location=device,
+)
+logit_lens = LogitLens.from_model(model)
 
 lens_options_dict = {
     "Tuned Lens": tuned_lens,
@@ -20,13 +23,15 @@ lens_options_dict = {
 
 statistic_options_dict = {
     "Entropy": "entropy",
-    "Cross Entropy": "ce",
+    "Cross Entropy": "cross_entropy",
     "Forward KL": "forward_kl",
 }
 
 
 def make_plot(lens, text, statistic, token_cutoff):
     input_ids = tokenizer.encode(text, return_tensors="pt")
+    input_ids = [tokenizer.bos_token_id] + input_ids
+    targets = input_ids[1:] + [tokenizer.eos_token_id]
 
     if len(input_ids[0]) == 0:
         return go.Figure(layout=dict(title="Please enter some text."))
@@ -34,18 +39,19 @@ def make_plot(lens, text, statistic, token_cutoff):
     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=2,
+    start_pos=max(len(input_ids[0]) - token_cutoff, 0),
+    pred_traj = PredictionTrajectory.from_lens_and_model(
+        lens=lens,
+        model=model,
         input_ids=input_ids,
-        start_pos=max(len(input_ids[0]) - token_cutoff, 0),
-        statistic=statistic_options_dict[statistic],
+        tokenizer=tokenizer,
+        targets=targets,
+        start_pos=start_pos,
     )
 
-    return fig
-
+    return getattr(pred_traj, statistic)().figure(
+        title=f"{lens.__class__.__name__} ({model.name_or_path}) {statistic}",
+    )
 
 preamble = """
 # The Tuned Lens 🔎

requirements.txt

@@ -1,2 +1,2 @@
-tuned_lens==0.0.5
+tuned_lens==0.1.0
 gradio