Update app.py

app.py

@@ -50,9 +50,9 @@ def get_hidden_states(raw_original_prompt):
     outputs = model(**model_inputs, output_hidden_states=True, return_dict=True)
     hidden_states = torch.stack([h.squeeze(0).cpu().detach() for h in outputs.hidden_states], dim=0)
     with gr.Row() as tokens_container:
-        for token in tokens:
-            gr.Button(token)
-    return tokens_container
+        # for token in tokens:
+        #     gr.Button(token)
+    return str(tokens)
 
 
 def run_model(raw_original_prompt, raw_interpretation_prompt, max_new_tokens, do_sample, 
@@ -115,8 +115,6 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                 This idea was explored in the paper **Patchscopes** ([Ghandeharioun et al., 2024](https://arxiv.org/abs/2401.06102)) and was later investigated further in **SelfIE** ([Chen et al., 2024](https://arxiv.org/abs/2403.10949)). 
                 An honorary mention for **Speaking Probes** ([Dar, 2023](https://towardsdatascience.com/speaking-probes-self-interpreting-models-7a3dc6cb33d6) -- my post!! 🥳)  which was a less mature approach but with the same idea in mind. 
                 We follow the SelfIE implementation in this space for concreteness. Patchscopes are so general that they encompass many other interpretation techniques too!!! 
-
-
                 
                 👾 **The idea is really simple: models are able to understand their own hidden states by nature!** 👾  
                 If I give a model a prompt of the form ``User: [X] Assistant: Sure'll I'll repeat your message`` and replace ``[X]`` *during computation* with the hidden state we want to understand, 
@@ -146,7 +144,7 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
     
     with gr.Group('Output'):
-        tokens_container = gr.Row() 
+        tokens_container = gr.Text() 
         with gr.Column() as interpretations_container:
             pass