Update app.py

app.py

@@ -9,6 +9,9 @@ import spaces
 import gradio as gr
 from esm_scripts.extract import run_demo
 from esm import pretrained, FastaBatchedDataset
+from data.evaluate_data.utils import Ontology
+import difflib
+import re
 
 
 # Load the model
@@ -20,6 +23,21 @@ model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
 model_esm.to('cuda')
 model_esm.eval()
 
+godb = Ontology(f'data/go1.4-basic.obo', with_rels=True)
+go_des = pd.read_csv('data/go_descriptions1.4.txt', sep='|', header=None)
+go_des.columns = ['id', 'text']
+go_des = go_des.dropna()
+go_des['id'] = go_des['id'].apply(lambda x: re.sub('_', ':', x))
+go_obo_set = set(go_des['id'].tolist())
+go_des['text'] = go_des['text'].apply(lambda x: x.lower())
+GO_dict = dict(zip(go_des['text'], go_des['id']))
+Func_dict = dict(zip(go_des['id'], go_des['text']))
+
+# terms_mf = pd.read_pickle('/cluster/home/wenkai/deepgo2/data/mf/terms.pkl')
+terms_mf = pd.read_pickle('data/terms/mf_terms.pkl')
+choices_mf = [Func_dict[i] for i in list(set(terms_mf['gos']))]
+choices = {x.lower(): x for x in choices_mf}
+
 
 @spaces.GPU
 def generate_caption(protein, prompt):
@@ -106,7 +124,16 @@ def generate_caption(protein, prompt):
     prediction = model.generate(samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1.,
                                 repetition_penalty=1.0)
 
-    return prediction
+    x = prediction[0]
+    x = [eval(i) for i in x.split('; ')]
+    pred_terms = []
+    for i in x:
+        txt = i[0]
+        prob = i[1]
+        sim_list = difflib.get_close_matches(txt.lower(), choices, n=1, cutoff=0.9)
+        if len(sim_list) > 0:
+            pred_terms.append((sim_list[0], prob))
+    return str(pred_terms)
     # return "test"