fix: add imports

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from speaking_probes.generate import extract_gpt_parameters, speaking_probe
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 @st.cache
@@ -16,7 +17,7 @@ model, model_params, tokenizer = load_model(model_name)
 
 neuron_layer = st.text_input("Layer: ")
 neuron_dim = st.text_dim("Dim: ")
-neurons = K_heads[int(neuron_layer), int(neuron_dim)]
+neurons = model_params.K_heads[int(neuron_layer), int(neuron_dim)]
 prompt = st.text_area("Prompt: ")
 submitted = st.button("Send!")