Update demo/modules/compute_score.py

demo/modules/compute_score.py

@@ -48,20 +48,14 @@ def compute_score(input_type_1: str, input_1: str, input_type_2: str, input_2: s
     return f"{score.item():.4f}"
 
 
-# Parse the uploaded structure file and return the sequence
-def pdb2seq(file):
-    parsed_seqs = get_struc_seq("/sujin/bin/foldseek", file)
-    
-    for seqs in parsed_seqs.values():
-        return seqs[0]
-
-
-# Parse the uploaded structure file and return the foldseek sequence
-def pdb2foldseek(file):
-    parsed_seqs = get_struc_seq("/sujin/bin/foldseek", file)
-    
+# Convert pdb file to aa sequence or foldseek sequence
+def parse_pdb(file, input_type):
+    parsed_seqs = get_struc_seq("/tmp/foldseek", file)
     for seqs in parsed_seqs.values():
-        return seqs[1].lower()
+        if input_type == "protein sequence":
+            return seqs[0]
+        else:
+            return seqs[1].lower()
 
 
 def change_input_type(choice_1: str, choice_2: str):
@@ -106,7 +100,7 @@ def build_score_computation():
                 
                 # Provide an upload button to upload a pdb file
                 upload_btn_1 = gr.UploadButton(label="Upload .pdb/.cif file", scale=0)
-                upload_btn_1.upload(pdb2seq, inputs=[upload_btn_1], outputs=[input_1])
+                upload_btn_1.upload(parse_pdb, inputs=[upload_btn_1, input_type_1], outputs=[input_1])
             
             with gr.Row():
                 input_2 = gr.Textbox(label="Input 2")
@@ -117,7 +111,7 @@ def build_score_computation():
                 
                 # Provide an upload button to upload a pdb file
                 upload_btn_2 = gr.UploadButton(label="Upload .pdb/.cif file", scale=0, visible=False)
-                upload_btn_2.upload(pdb2seq, inputs=[upload_btn_2], outputs=[input_2])
+                upload_btn_2.upload(parse_pdb, inputs=[upload_btn_2, input_type_2], outputs=[input_2])
             
             # Provide examples
             examples = gr.Dataset(samples=samples, type="index", components=[input_1, input_2], label="Input examples")