use Gante GPT2 code

app.py

@@ -1,16 +1,14 @@
 import gradio as gr
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import GPT2Tokenizer, AutoModelForCausalLM
 import numpy as np
 
-
 MODEL_NAME = "gpt2"
 
-
 if __name__ == "__main__":
     # Define your model and your tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)  # or AutoModelForCausalLM
+    tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
     if tokenizer.pad_token_id is None:
         tokenizer.pad_token_id = tokenizer.eos_token_id
         model.config.pad_token_id = model.config.eos_token_id
@@ -28,70 +26,54 @@ if __name__ == "__main__":
         "p < 1%": "red"
     }
 
+
     def get_tokens_and_labels(prompt):
         """
         Given the prompt (text), return a list of tuples (decoded_token, label)
         """
         inputs = tokenizer([prompt], return_tensors="pt")
         outputs = model.generate(
-            **inputs, max_new_tokens=50, return_dict_in_generate=True, output_scores=True
+            **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)
         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_ids = outputs.sequences[:, input_length:]
-        generated_tokens = tokenizer.convert_ids_to_tokens(generated_ids[0])
+        generated_tokens = outputs.sequences[:, input_length:]
 
-        # Important: you might need to find a tokenization character to replace (e.g. "Ġ" for BPE) and get the correct
-        # spacing into the final output 👼
-        if model.config.is_encoder_decoder:
-            highlighted_out = []
-        else:
-            input_tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids)
-            highlighted_out = [(token.replace("▁", " "), None) for token in input_tokens]
+        # 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
-        for token, proba in zip(generated_tokens, transition_proba[0]):
+        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((token.replace("▁", " "), this_label))
+            highlighted_out.append((tokenizer.decode(token), this_label))
 
         return highlighted_out
 
+
     demo = gr.Blocks()
     with demo:
         gr.Markdown(
             """
-            # 🌈 Color-Coded Text Generation 🌈
-
+            # 🌈 Color Coded Text Generation 🌈
             This is a demo of how you can obtain the probabilities of each generated token, and use them to
-            color code the model output. Internally, it relies on
-            [`compute_transition_scores`](https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores),
+            color code the model output.
+            Feel free to clone this demo and modify it to your needs 🤗
+            Internally, it relies on [`compute_transition_scores`](https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores),
             which was added in `transformers` v4.26.0.
-
-            ⚠️ For instance, with the pre-populated input and its color-coded output, you can see that
-            `google/flan-t5-base` struggles with arithmetics.
-
-            🤗 Feel free to clone this demo and modify it to your needs 🤗
             """
         )
 
         with gr.Row():
             with gr.Column():
-                prompt = gr.Textbox(
-                    label="Prompt",
-                    lines=3,
-                    value=(
-                        "Answer the following question by reasoning step-by-step. The cafeteria had 23 apples. "
-                        "If they used 20 for lunch and bought 6 more, how many apples do they have?"
-                    ),
-                )
-                button = gr.Button(f"Generate with {MODEL_NAME}")
+                prompt = gr.Textbox(label="Prompt", lines=3, value="Today is")
+                button = gr.Button(f"Generate with {MODEL_NAME}, using sampling!")
             with gr.Column():
                 highlighted_text = gr.HighlightedText(
                     label="Highlighted generation",
@@ -101,6 +83,6 @@ if __name__ == "__main__":
 
         button.click(get_tokens_and_labels, inputs=prompt, outputs=highlighted_text)
 
-
 if __name__ == "__main__":
     demo.launch()
+