Update app.py

app.py

@@ -49,6 +49,9 @@ def generate_text_image_with_pil(old_text, input_token, new_token, image_width=4
     paragraphs = old_text.split('\n')
     
     # Add the token information to the last paragraph
+    input_token = input_token.replace("\n","\\n")
+    new_token = new_token.replace("\n","\\n")
+    
     if paragraphs:
         paragraphs[-1] += f"[{input_token}]→[{new_token}]"
     else:
@@ -72,7 +75,7 @@ def generate_text_image_with_pil(old_text, input_token, new_token, image_width=4
     draw = ImageDraw.Draw(img)
     
     # Load font
-    font_path = "NotoSansCJK-Bold.ttc"
+    font_path = "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc"
     font = ImageFont.truetype(font_path, font_size)
     
     # Draw text
@@ -80,7 +83,7 @@ def generate_text_image_with_pil(old_text, input_token, new_token, image_width=4
     token_marker = f"[{input_token}]→[{new_token}]"
     
     for line in all_lines:
-        if token_marker in line:
+        if token_marker in line:                
             parts = line.split(token_marker)
             # Draw text before token
             draw.text((10, y), parts[0], fill="black", font=font)
@@ -102,11 +105,12 @@ def generate_text_image_with_pil(old_text, input_token, new_token, image_width=4
                 x += draw.textlength(f"[{new_token}]", font=font)
                 draw.text((x, y), parts[1], fill="black", font=font)
         else:
+            print(token_marker)
+            print(line)
             draw.text((10, y), line, fill="black", font=font)
         
         # Move to next line, adding extra space between paragraphs
         y += font_size + 8
-    
     return np.array(img)
 
 
@@ -163,7 +167,10 @@ def render_next_token_table_image(table_data, predict_token, image_width=500, ro
                     token_str = token_str
                     color = "red" if token_str == predict_token and col == 1 else "blue" if col == 1 else "black"
                     bold = token_str == predict_token and col == 1
-                    token_str_ = token_str.replace("\n", "\\n").replace(" ", "\\s").replace("\t", "\\t")
+                    if token_str.count(" ") ==  1 and len(token_str) != 1:
+                        token_str_ = token_str.replace("\n", "\\n").replace("\t", "\\t")
+                    else:
+                        token_str_ = token_str.replace("\n", "\\n").replace(" ", "\\s").replace("\t", "\\t")
                     draw_cell(x, y, f"{token_str_} ({prob:.1%})", color=color, bold=bold)
 
     return np.array(img)