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()