columns and printing fixed

app.py

@@ -15,15 +15,20 @@ def load_model(model_name):
 model_name = st.selectbox("Select a model: ", options=['gpt2', 'gpt2-medium', 'gpt2-large'])
 model, model_params, tokenizer = load_model(model_name)
 
-neuron_layer = st.text_input("Layer: ")
-neuron_dim = st.text_input("Dim: ")
+col1, col2, *_ = st.columns(5)
+neuron_layer = col1.text_input("Layer: ", value='0')
+neuron_dim = col2.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:
-    speaking_probe(model, model_params, tokenizer, prompt, *neurons, num_generations=1, 
+    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)
+
+    for text in decoded:
+        st.code(text, language=None)

speaking_probes/generate.py

@@ -151,7 +151,6 @@ def speaking_probe(model, model_params, tokenizer, prompt, *neurons,
     if len(bad_words_ids) == 0:
         bad_words_ids = None
         
-    print(prompt)
     input_ids = tokenizer_with_neurons.encode(prompt, return_tensors='pt').to(model.device)
     input_ids = torch.cat([deepcopy(input_ids) for _ in range(num_generations)], dim=0)
     outputs = model.generate(input_ids, pad_token_id=model.config.eos_token_id, 
@@ -161,18 +160,17 @@ def speaking_probe(model, model_params, tokenizer, prompt, *neurons,
                              **kwargs)
 
     decoded = tokenizer_with_neurons.batch_decode(outputs.sequences, skip_special_tokens=True)
-    
-    for i in range(len(decoded)):
-        print("\n\ngenerate:", decoded[i])
-    
+        
     if has_extra_neurons:
         model.resize_token_embeddings(num_non_neuron_tokens)
         model.transformer.wte.weight.data = model.transformer.wte.weight.data[:num_non_neuron_tokens]
         
     if return_outputs:
-        return outputs
+        return decoded, outputs
+    else:
+        return decoded
+
 
-    
 # main
 if __name__ == "__main__":
     parser = ArgumentParser()
@@ -209,10 +207,13 @@ if __name__ == "__main__":
     i1, i2 = map(lambda x: int(x.strip()), args.neuron.split(','))
     neuron = model_params.K_heads[i1, i2]
     neurons = [neuron]
-    
-    speaking_probe(model, model_params, tokenizer, prompt, *neurons,
+
+    print(prompt)
+    decoded = speaking_probe(model, model_params, tokenizer, prompt, *neurons,
                    num_generations=args.num_generations, 
                    repetition_penalty=args.repetition_penalty, 
                    num_beams=args.num_beams, top_p=args.top_p, top_k=args.top_k,
                    min_length=args.min_length, do_sample=not args.no_sample,
                    max_length=args.max_length, max_new_tokens=args.max_new_tokens)
+    for i in range(len(decoded)):
+        print("\n\ngenerate:", decoded[i])