working local gradio server

api_prediction.py

@@ -143,7 +143,7 @@ class AptaTransPipeline_Dist(AptaTransPipeline):
         print('Predicting the Aptamer-Protein Interaction')
         try:
             print("loading the best model for api!")
-            self.model.load_state_dict(torch.load('./models/test_lr=1e-06_batch_size=16_epochs=30.pt', map_location=self.device))
+            self.model.load_state_dict(torch.load('./models/model.pt', map_location=self.device))
         except:
             print('there is no best model file.')
             print('You need to train the model for predicting API!')

inference.py

@@ -1,13 +1,8 @@
 from api_prediction import AptaTransPipeline_Dist
 import argparse
+import gradio as gr
 
 def infer(prot, apta):
-
-    print('Input protein: ', prot)
-    print('Input aptamer: ', apta)
-
-    print('Initializing model.')
-
     pipeline = AptaTransPipeline_Dist(
         lr=1e-6,
         weight_decay=None,
@@ -25,20 +20,29 @@ def infer(prot, apta):
         load_best_pt=True, # already loads the pretrained model using the datasets included in repo -- no need to run the bottom two cells
         device='cuda',
         seed=1004)
+    
     scores = pipeline.inference([apta], [prot], [0])
-    print('Your predicted score is: ', scores[0])
+    return scores[0]
+
+def comparison(protein, aptamers):
+    aptamers = aptamers.split('\n')
+    pairs = [[protein, aptamer] for aptamer in aptamers]
+    print(pairs)
+    scores = []
+    
+    for pair in pairs:
+        score = infer(pair[0], pair[1])
+        scores.append(score)
 
-def main():
-    parser = argparse.ArgumentParser(
-                        prog='API Pipeline Inference',
-                        description='From a protein and RNA aptamer sequence, predict the binding score.')
-    parser.add_argument('-p', '--prot')
-    parser.add_argument('-a', '--apta')
+    return scores
 
-    args = parser.parse_args()
-    protein = str(args.prot)
-    apta = str(args.apta)
-    infer(protein, apta)
+iface = gr.Interface(
+    fn=comparison,
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Protein"),
+        gr.Textbox(lines=10,placeholder="Aptamers (1 per line)")
+    ],
+    outputs=gr.Textbox()
+)
 
-if __name__ == "__main__":
-    main()
+iface.launch()