added colors

app.py

@@ -32,10 +32,17 @@ if __name__ == "__main__":
     ]
 
     label_to_color = {
-        "p >= 80%": "blue",
-        "p >= 50%": "green",
-        "p >= 20%": "yellow",
-        "p >= 0%": "red"
+        "p >= 80%": "#d91111",
+        "p >= 90%": "#d97211",
+        "p >= 80%": "#d99a11",
+        "p >= 70%": "#d9c111",
+        "p >= 60%": "#d5d911",
+        "p >= 50%": "#a0d911",
+        "p >= 40%": "#4dd911",
+        "p >= 30%": "#11d954",
+        "p >= 20%": "#11d9a0",
+        "p >= 10%": "#11d9d2",
+        "p >= 00%": "#11b4d9"
     }
 
     def get_tokens_and_labels(prompt):
@@ -48,7 +55,7 @@ if __name__ == "__main__":
         )
         # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
         transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
-        transition_proba = np.exp(transition_scores[0])
+        transition_proba = np.exp(transition_scores)
         # We only have scores for the generated tokens, so pop out the prompt tokens
         input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
         generated_ids = outputs.sequences[:, input_length:]