1st commit

app.py

@@ -0,0 +1,62 @@
+import gradio as gr
+
+from transformers import GPT2Tokenizer, AutoModelForCausalLM
+import numpy as np
+
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+tokenizer.pad_token_id = tokenizer.eos_token_id
+
+# if prob > x, then label = y; sorted in descending probability order
+probs_to_label = [
+    (0.1, "p >= 10%"),
+    (0.01, "p >= 1%"),
+    (1e-20, "p < 1%"),
+]
+
+label_to_color = {
+    "p >= 10%": "green",
+    "p >= 1%": "yellow",
+    "p < 1%": "red"
+}
+
+def get_tokens_and_scores(prompt):
+    inputs = tokenizer([prompt], return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=50, return_dict_in_generate=True, output_scores=True, do_sample=True)
+    transition_scores = model.compute_transition_scores(
+        outputs.sequences, outputs.scores, normalize_logits=True
+    )
+    transition_proba = np.exp(transition_scores)
+    input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
+    generated_tokens = outputs.sequences[:, input_length:]
+    highlighted_out = [(tokenizer.decode(token), None) for token in inputs.input_ids]
+
+    for token, proba in zip(generated_tokens[0], transition_proba[0]):
+        this_label = None
+        assert 0. <= proba <= 1.0
+        for min_proba, label in probs_to_label:
+            if proba >= min_proba:
+                this_label = label
+                break
+        highlighted_out.append((tokenizer.decode(token), this_label))
+
+    return highlighted_out
+
+
+demo = gr.Interface(
+    get_tokens_and_scores,
+    [
+        gr.Textbox(
+          label="Prompt",
+          lines=3,
+          value="Today is",
+        ),
+    ],
+    gr.HighlightedText(
+        label="Highlighted generation",
+        combine_adjacent=True,
+        show_legend=True,
+    ).style(color_map=label_to_color),
+)
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1 @@
+transformers>=4.26