add more viz

app.py

@@ -7,6 +7,7 @@ from transformers import AutoTokenizer
 
 tokenizer = AutoTokenizer.from_pretrained("adamcasson/ul2-tinystories")
 
+
 def mask_spans(
     tokens,
     mu,
@@ -19,7 +20,9 @@ def mask_spans(
     masked_tokens = tokens[:]
 
     encoder_inputs = [prepend_id] if prepend_id is not None else []
+    encoder_mask = [1] if prepend_id is not None else []
     targets = []
+    targets_mask = []
 
     # Original T5 code reused tokens at the end of vocab for sentinels
     # https://github.com/google-research/text-to-text-transfer-transformer/blob/258fd30687e6c60d18b7204d009dc5c753142987/t5/data/preprocessors.py#L3106C6-L3106C6
@@ -32,7 +35,9 @@ def mask_spans(
             0, len(tokens) - random.randint(1, int(2 * mu))
         )  # max to handle start < 0
         encoder_inputs += tokens[:start] + [sentinel_id]
-        targets += tokens[start:]
+        encoder_mask += ([1] * len(tokens[:start])) + [0]
+        targets += [sentinel_id] + tokens[start:]
+        targets_mask += [0] + ([1] * len(tokens[start:]))
         for i in range(start, len(tokens)):
             masked_tokens[i] = -1
 
@@ -49,28 +54,30 @@ def mask_spans(
             # randomly decide if span should be masked
             if np.random.binomial(1, p=r):
                 encoder_inputs.append(sentinel_id)
+                encoder_mask.append(0)
                 targets += tokens[start:end]
+                targets_mask += ([1] * len(tokens[start:end]))
                 for i in range(start, end):
                     masked_tokens[i] = -1
                 prev_span_unmasked = False
                 sentinel_id -= 1
             else:
                 encoder_inputs += tokens[start:end]
+                encoder_mask += ([1] * len(tokens[start:end]))
                 # if previous span was also unmasked we don't need to keep adding the sentinel token
                 if not prev_span_unmasked:
                     targets.append(sentinel_id)
+                    targets_mask.append(0)
                     prev_span_unmasked = True
             start = end
 
-    encoder_inputs.append(eos_id)
     targets.append(eos_id)
-    decoder_inputs = (
-        [prepend_id] + targets[:-1]
-        if prepend_id is not None
-        else [eos_id] + targets[:-1]
-    )
+    targets_mask.append(1)
+    decoder_inputs = [eos_id] + targets[:-1]
+    decoder_mask = [1] + targets_mask[:-1]
+
+    return encoder_inputs, encoder_mask, decoder_inputs, decoder_mask, targets, targets_mask, masked_tokens
 
-    return encoder_inputs, decoder_inputs, targets, masked_tokens
 
 # Create mixture-of-denoisers
 denoiser_map = {
@@ -137,15 +144,23 @@ def mask_viz(denoiser, text):
     seq = tokenizer.encode(text)
     tokens = tokenizer.tokenize(text)
 
-    out = denoiser_map[denoiser](seq)
+    enc_in, enc_mask, dec_in, dec_mask, targets, targets_mask, mask = denoiser_map[denoiser](seq)
     
-    mask = out[-1]
-
     highlight_tok = []
     for tok, tok_mask in zip(tokens, mask):
         highlight_tok.append((tok.replace("Ġ", " ").replace("Ċ", "\n"), "masked" if tok_mask == -1 else "unmasked"))
 
-    return highlight_tok
+    highlight_enc = []
+    enc_tok = tokenizer.convert_ids_to_tokens(enc_in)
+    for id, tok, tok_mask in zip(enc_in, enc_tok, enc_mask):
+        highlight_enc.append((tok.replace("Ġ", " ").replace("Ċ", "\n") if tok_mask == 1 else str(id), "masked" if tok_mask == 0 else "unmasked"))
+
+    highlight_dec = []
+    dec_tok = tokenizer.convert_ids_to_tokens(dec_in)
+    for id, tok, tok_mask in zip(dec_in, dec_tok, dec_mask):
+        highlight_dec.append((tok.replace("Ġ", " ").replace("Ċ", "\n") if tok_mask == 1 else str(id), "masked" if tok_mask == 0 else "unmasked"))
+
+    return highlight_tok, highlight_enc, highlight_dec
 
 iface = gr.Interface(
     fn=mask_viz,
@@ -164,14 +179,29 @@ iface = gr.Interface(
             value="R (µ = 3, r = 0.15)",
         ),
         gr.Textbox(
-            value='Once upon a time, there was a clever little dog named Max. Max loved to run and play with his friends in the park. One day, Max was running very fast when he fell and hurt his knee. Max went to his friend, the wise old owl, and said, "Owl, my knee hurts. What can I do?" The owl thought for a moment and said, "Max, you should test your knee. Try to walk slowly and see if it still hurts." So Max tested his knee by walking slowly. At first, it hurt a little, but soon Max felt better. He said, "Thank you, Owl, for your help. Now I can play with my friends again." Max was so happy that he could play with his friends without pain. He learned that sometimes, it was good to slow down and listen to his body. And Max and his friends played happily in the park ever after.'
+            value='Once upon a time, there was a family with a little boy. His name was Jack.\nOne day, Jack had a thought. He wanted to go to the park and play. His parents were worried because it was getting dark and the park was far away.\n"Mom, I want to play in the park," Jack said.\nHis mother thought for a moment. "It\'s too late to go to the park now. We\'d better stay at home," she said. \nJack was sad, but he understood why his parents were worried. Together they decided to play games at home instead. \nJack was so happy to get to play games with his family. He thought it was the best time ever.'
+        ),
+    ],
+    outputs=[
+        gr.HighlightedText(
+            label="Corrupted spans",
+            combine_adjacent=True,
+            show_legend=True,
+            color_map={"unmasked": "green", "masked": "red"}
+        ),
+        gr.HighlightedText(
+            label="Encoder input",
+            combine_adjacent=True,
+            show_legend=True,
+            color_map={"unmasked": "green", "masked": "red"}
+        ),
+        gr.HighlightedText(
+            label="Decoder input",
+            combine_adjacent=True,
+            show_legend=True,
+            color_map={"unmasked": "green", "masked": "red"}
         ),
     ],
-    outputs=gr.HighlightedText(
-        combine_adjacent=True,
-        show_legend=True,
-        color_map={"unmasked": "green", "masked": "red"}
-    )
 )
 
 iface.launch()