improve style and avoid cache mutation

app.py

@@ -12,23 +12,22 @@ def load_model(model_name):
     return model, model_params, tokenizer
     
 
-model_name = st.selectbox("Select a model: ", options=['gpt2', 'gpt2-medium', 'gpt2-large'])
+col1, col2, col3, *_ = st.columns(5)
+model_name = col1.selectbox("Select a model: ", options=['gpt2', 'gpt2-medium', 'gpt2-large'])
 model, model_params, tokenizer = load_model(model_name)
-
-col1, col2, *_ = st.columns(5)
-neuron_layer = col1.text_input("Layer: ", value='0')
-neuron_dim = col2.text_input("Dim: ", value='0')
+neuron_layer = col2.text_input("Layer: ", value='0')
+neuron_dim = col3.text_input("Dim: ", value='0')
 
 neurons = model_params.K_heads[int(neuron_layer), int(neuron_dim)]
 prompt = st.text_area("Prompt: ")
 submitted = st.button("Send!")
 
-
 if submitted:
+    model, model_params, tokenizer = map(deepcopy, (model, model_params, tokenizer))
     decoded = speaking_probe(model, model_params, tokenizer, prompt, *neurons, num_generations=1, 
-                   repetition_penalty=2., 
-                   num_beams=3, min_length=1, do_sample=True, 
-                   max_new_tokens=100)
+                             repetition_penalty=2., num_generations=3,
+                             min_length=1, do_sample=True, 
+                             max_new_tokens=100)
 
     for text in decoded:
-        st.code(text, language=None)
+        st.text(text)