derp

app.py

@@ -36,16 +36,16 @@ if __name__ == "__main__":
         outputs = model.generate(
             **inputs, max_new_tokens=50, return_dict_in_generate=True, output_scores=True, do_sample=True
         )
-        # Important, don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
+        # 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)
         # 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_tokens = outputs.sequences[:, input_length:]
 
-        # initialize the highlighted output with the prompt, which will have no color label
+        # Initialize the highlighted output with the prompt, which will have no color label
         highlighted_out = [(tokenizer.decode(token), None) for token in inputs.input_ids]
-        # get the (decoded_token, label) pairs for the generated tokens
+        # Get the (decoded_token, label) pairs for the generated tokens
         for token, proba in zip(generated_tokens[0], transition_proba[0]):
             this_label = None
             assert 0. <= proba <= 1.0
@@ -70,7 +70,7 @@ if __name__ == "__main__":
             label="Highlighted generation",
             combine_adjacent=True,
             show_legend=True,
-        ).style(color_map=label_to_color),
+        ).style(color_map=label_to_color)
         button = gr.Button(f"Generate with {MODEL_NAME}")
 
         button.click(get_tokens_and_labels, inputs=prompt, outputs=highlighted_text)