fixed several bugs in app.py

.dockerignore

@@ -1,2 +1,2 @@
 lens
-.git
+.git

README.md

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

app.py

@@ -7,7 +7,7 @@ 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")
+model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-410m-deduped", torch_dtype="auto")
 model = model.to(device)
 tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-410m-deduped")
 tuned_lens = TunedLens.from_model_and_pretrained(
@@ -29,19 +29,19 @@ statistic_options_dict = {
 
 
 def make_plot(lens, text, statistic, token_cutoff):
-    input_ids = tokenizer.encode(text, return_tensors="pt")
+    input_ids = tokenizer.encode(text)
     input_ids = [tokenizer.bos_token_id] + input_ids
     targets = input_ids[1:] + [tokenizer.eos_token_id]
 
-    if len(input_ids[0]) == 0:
+    if len(input_ids) == 1:
         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."))
 
-    start_pos=max(len(input_ids[0]) - token_cutoff, 0),
+    start_pos=max(len(input_ids) - token_cutoff, 0)
     pred_traj = PredictionTrajectory.from_lens_and_model(
-        lens=lens,
+        lens=lens_options_dict[lens],
         model=model,
         input_ids=input_ids,
         tokenizer=tokenizer,
@@ -49,7 +49,7 @@ def make_plot(lens, text, statistic, token_cutoff):
         start_pos=start_pos,
     )
 
-    return getattr(pred_traj, statistic)().figure(
+    return getattr(pred_traj, statistic_options_dict[statistic])().figure(
         title=f"{lens.__class__.__name__} ({model.name_or_path}) {statistic}",
     )
 
@@ -114,5 +114,4 @@ with gr.Blocks() as demo:
     demo.load(make_plot, [lens_options, text, statistic, token_cutoff], plot)
 
 if __name__ == "__main__":
-    demo.launch()
-
+    demo.launch(server_name="0.0.0.0", server_port=7860)