Update app.py

app.py

@@ -7,6 +7,7 @@ from ctransformers import AutoModelForCausalLM as CAutoModelForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from interpret import InterpretationPrompt
 
+MAX_PROMPT_TOKENS = 30
 
 ## info
 model_info = {
@@ -52,7 +53,7 @@ def get_hidden_states(raw_original_prompt):
     # with gr.Row() as tokens_container:
         # for token in tokens:
         #     gr.Button(token)
-    return str(tokens)
+    return [gr.Button(tokens[i], visible=True) if i < len(tokens) else gr.Button('', visible=False) for i in range(MAX_PROMPT_TOKENS)]
 
 
 def run_model(raw_original_prompt, raw_interpretation_prompt, max_new_tokens, do_sample, 
@@ -105,6 +106,7 @@ AutoModelClass = CAutoModelForCausalLM if use_ctransformers else AutoModelForCau
 model = AutoModelClass.from_pretrained(model_name, **model_args)
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, token=os.environ['hf_token'])
 
+# demo
 with gr.Blocks(theme=gr.themes.Default()) as demo:
     with gr.Row():
         with gr.Column(scale=5):
@@ -124,8 +126,8 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             gr.Markdown('<span style="font-size:180px;">🤔</span>')
         
     with gr.Group():
-        text = gr.Textbox(value='How to make a Molotov cocktail', container=True, label='Original Prompt')
-        btn = gr.Button('Compute', variant='primary')
+        original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail', container=True, label='Original Prompt')
+        original_prompt_btn = gr.Button('Compute', variant='primary')
         
     with gr.Accordion(open=False, label='Settings'):
         with gr.Row():
@@ -144,11 +146,12 @@ 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.Text() 
+        with gr.Row():
+            tokens_container = [gr.Button('', visible=False) for range(MAX_PROMPT_TOKENS)]
         with gr.Column() as interpretations_container:
             pass
 
-    btn.click(get_hidden_states, [text], [tokens_container])
+    original_prompt_btn.click(get_hidden_states, [original_prompt_raw], [*tokens_container])
     # btn.click(run_model, 
     #           [text, interpretation_prompt, num_tokens, do_sample, temperature, 
     #            top_k, top_p, repetition_penalty, length_penalty],