Update app.py

app.py

@@ -10,14 +10,19 @@ from huggingface_hub import hf_hub_download
 
 def bold_predicted_letters(input_string: str, output_string: str) -> str:
     result = []
-    i = 0
+    i = j = 0
+    input_string = input_string.upper()
+    output_string = output_string.upper()
     while i < len(input_string):
-        if input_string[i:i+6] == "<mask>":
-            result.append("**" + output_string[i] + "**")
+        if input_string[i:i+6] == "<MASK>":
+            result.append("**" + output_string[j] + "**")
             i += 6
+            j += 1
         else:
-            result.append(output_string[i])
+            result.append(input_string[i])
             i += 1
+            if input_string[i-1] != "<":
+                j += 1
     return "".join(result)
 
 class model:
@@ -92,9 +97,6 @@ class model:
         formatted_predicted_sequence = formatted_predicted_sequence.replace("<pad>","")
         formatted_predicted_sequence = formatted_predicted_sequence.replace("<cls>","")
         formatted_predicted_sequence = formatted_predicted_sequence.replace("<eos>","")
-        
-        print(sequence_input)
-        print(formatted_predicted_sequence)
 
         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