Update app.py

app.py

@@ -8,6 +8,18 @@ from celle_main import instantiate_from_config
 from omegaconf import OmegaConf
 from huggingface_hub import hf_hub_download
 
+def bold_predicted_letters(input_string: str, output_string: str) -> str:
+    result = []
+    i = 0
+    while i < len(input_string):
+        if input_string[i:i+6] == "<mask>":
+            result.append("**" + output_string[i] + "**")
+            i += 6
+        else:
+            result.append(output_string[i])
+            i += 1
+    return "".join(result)
+
 class model:
     def __init__(self):
         self.model = None
@@ -75,8 +87,13 @@ class model:
             model=self.model,
             device=device,
         )
+
+        formatted_predicted_sequence = formatted_predicted_sequence.replace("<pad>","")
+        formatted_predicted_sequence = formatted_predicted_sequence("<start>","")
+        formatted_predicted_sequence = formatted_predicted_sequence("<end>","")
+        formatted_predicted_sequence = bold_predicted_letters(sequence_input, formatted_predicted_sequence)                                                      
         
-        return T.ToPILImage()(protein_image[0,0]), T.ToPILImage()(nucleus_image[0,0]), formatted_predicted_sequence.replace("<pad>", "")
+        return T.ToPILImage()(protein_image[0,0]), T.ToPILImage()(nucleus_image[0,0]), 
 
 base_class = model()