Update app.py

app.py

@@ -17,7 +17,7 @@ dataset = load_dataset('dar-tau/grammar-attention-maps-opt-350m')['train']
 tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m', add_prefix_space=True)
 
 
-def analyze_sentence(index, vis_type):
+def analyze_sentence(index, vis_type, vis_format):
     row = dataset[index]
     text = row['text']
     tokenized = tokenizer.batch_decode(tokenizer.encode(text, add_special_tokens=False))
@@ -55,7 +55,8 @@ with demo:
         btn = gr.Button("Run", min_width=30)
     output = gr.Plot(label="Plot", container=True)
     metrics = gr.Label("Metrics")
-    btn.click(analyze_sentence, [sentence_dropdown, vis_dropdown], [output, metrics])
+    vis_format_checkbox = gr.Radio(['Plot', 'Graph'])
+    btn.click(analyze_sentence, [sentence_dropdown, vis_dropdown, vis_format_checkbox], [output, metrics])
 
 
 if __name__ == "__main__":