Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import os
 import gc
+import json
 
 import numpy as np
 import pandas as pd
@@ -97,11 +98,18 @@ def get_model():
 
 
 model = get_model()
-
-def greet(seq):
-    processed_seq = Sequence.prepare(seq)
-    raw_prediction = model.predict(processed_seq)
-    return f"Input is {seq}.\nProcessed input is:\n{processed_seq}\n\nModel makes prediction:\n{raw_prediction}"
+mappings_path = cached_download(hf_hub_url("jonathang/Protein_Family_CNN", 'prot_mappings.json'))
+with open(mappings_path) as f:
+    prot_mappings = json.load(f)
+
+def greet(Amino_Acid_Sequence):
+    processed_seq = Sequence.prepare(Amino_Acid_Sequence)
+    raw_prediction = model.predict(processed_seq)[0]
+    idx = raw_prediction.argmax()
+    fam_asc = prot_mappings['id2fam_asc'][idx]
+    fam_id = prot_mappings['fam_asc2fam_id'][fam_asc]
+    gc.collect()
+    return f"Input is {Amino_Acid_Sequence}.\nProcessed input is:\n{processed_seq}\n\nModel makes prediction:\n{raw_prediction}\nWhich gives index {idx}.\nWhich maps to family ascension {fam_asc} and ID {fam_id}"
 
 iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 iface.launch()