Update app.py

app.py

@@ -7,42 +7,31 @@ from transformers import AutoTokenizer, AutoModel
 from huggingface_hub import hf_hub_download
 from keras.models import load_model
 
-# ———————————————————  CONFIGURAÇÃO  ——————————————————— #
-SPACE_ID      = "melvinalves/protein_function_prediction"
-TOP_N         = 10
-THRESH        = 0.50                # limiar para listar GO terms
-CHUNK_PB      = 512
-CHUNK_ESM     = 1024
-
-# ———————————————————  HELPERS DE CACHE  ——————————————————— #
+# ———————————————————  CONFIG  ——————————————————— #
+SPACE_ID   = "melvinalves/protein_function_prediction"
+TOP_N      = 10
+THRESH     = 0.50
+CHUNK_PB   = 512
+CHUNK_ESM  = 1024
+
+# ———————————————————  HELPERS  ——————————————————— #
 @st.cache_resource
-def download_file(path_in_repo: str):
-    return hf_hub_download(repo_id=SPACE_ID, repo_type="space", filename=path_in_repo)
+def download_file(path):
+    return hf_hub_download(repo_id=SPACE_ID, repo_type="space", filename=path)
 
 @st.cache_resource
-def load_keras(file_name: str):
-    return load_model(download_file(f"models/{file_name}"), compile=False)
+def load_keras(name):
+    return load_model(download_file(f"models/{name}"), compile=False)
 
 @st.cache_resource
-def load_hf_encoder(model_name: str):
-    tok = AutoTokenizer.from_pretrained(model_name, do_lower_case=False)
-    mdl = AutoModel.from_pretrained(model_name)
+def load_hf_encoder(model):
+    tok = AutoTokenizer.from_pretrained(model, do_lower_case=False)
+    mdl = AutoModel.from_pretrained(model)
     mdl.eval()
     return tok, mdl
 
-# ———————————————————  MODELOS  ——————————————————— #
-mlp_pb   = load_keras("mlp_protbert.h5")
-mlp_bfd  = load_keras("mlp_protbertbfd.h5")
-mlp_esm  = load_keras("mlp_esm2.h5")
-stacking = load_keras("ensemble_stack.h5")      # usa o nome que tiveres guardado
-
-# ———————————————————  LABEL BINARIZER  ——————————————————— #
-mlb       = joblib.load(download_file("data/mlb_597.pkl"))
-GO_TERMS  = mlb.classes_
-
-# ———————————————————  EMBEDDINGS  ——————————————————— #
-def embed_seq(model_name: str, seq: str, chunk: int) -> np.ndarray:
-    tok, mdl = load_hf_encoder(model_name)
+def embed_seq(model, seq, chunk):
+    tok, mdl = load_hf_encoder(model)
     parts = [seq[i:i+chunk] for i in range(0, len(seq), chunk)]
     vecs  = []
     for p in parts:
@@ -51,6 +40,15 @@ def embed_seq(model_name: str, seq: str, chunk: int) -> np.ndarray:
         vecs.append(out.last_hidden_state[:, 0, :].squeeze().numpy())
     return np.mean(vecs, axis=0, keepdims=True)
 
+# ———————————————————  CARGA MODELOS  ——————————————————— #
+mlp_pb   = load_keras("mlp_protbert.h5")
+mlp_bfd  = load_keras("mlp_protbertbfd.h5")
+mlp_esm  = load_keras("mlp_esm2.h5")
+stacking = load_keras("ensemble_stack.h5")   # usa o nome real aqui
+
+mlb      = joblib.load(download_file("data/mlb_597.pkl"))
+GO       = mlb.classes_
+
 # ———————————————————  UI  ——————————————————— #
 st.title("🔬 Predição de Funções de Proteínas")
 
@@ -61,47 +59,42 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
-fasta = st.text_area("Insere a sequência FASTA:", height=200)
+fasta_input = st.text_area("Insere a sequência FASTA:", height=200)
+predict_clicked = st.button("Prever GO terms")
 
-# ---------- BOTÃO ----------
-if fasta and st.button("Prever GO terms"):
-    seq = "\n".join(l for l in fasta.splitlines() if not l.startswith(">"))
-    seq = seq.replace(" ", "").replace("\n", "").upper()
+if predict_clicked:
 
+    # ——— Validação mínima ———
+    seq = "\n".join(l for l in fasta_input.splitlines() if not l.startswith(">")).replace(" ", "").upper()
     if not seq:
-        st.warning("Por favor, insere uma sequência válida.")
+        st.warning("Por favor, insere primeiro uma sequência FASTA válida.")
         st.stop()
 
-    # 1) EMBEDDINGS
-    st.write("⏳ A gerar embeddings…")
-    emb_pb  = embed_seq("Rostlab/prot_bert",            seq, CHUNK_PB)
-    emb_bfd = embed_seq("Rostlab/prot_bert_bfd",        seq, CHUNK_PB)
-    emb_esm = embed_seq("facebook/esm2_t33_650M_UR50D", seq, CHUNK_ESM)
-
-    # 2) PREDIÇÕES INDIVIDUAIS
-    st.write("🧠 A fazer predições…")
-    y_pb  = mlp_pb.predict(emb_pb)
-    y_bfd = mlp_bfd.predict(emb_bfd)
-    y_esm = mlp_esm.predict(emb_esm)[:, :597]      # corta 602 → 597
-
-    # 3) ENSEMBLE
-    X_stack = np.concatenate([y_pb, y_bfd, y_esm], axis=1)   # (1, 1791)
-    y_ens   = stacking.predict(X_stack)
-
-    # ——— Função auxiliar para mostrar resultados ———
-    def show_results(label: str, y_pred):
-        with st.expander(label, expanded=(label == "Ensemble (Stacking)")):
+    # ——— 1) EMBEDDINGS ———
+    with st.spinner("⏳ A gerar embeddings…"):
+        emb_pb  = embed_seq("Rostlab/prot_bert",            seq, CHUNK_PB)
+        emb_bfd = embed_seq("Rostlab/prot_bert_bfd",        seq, CHUNK_PB)
+        emb_esm = embed_seq("facebook/esm2_t33_650M_UR50D", seq, CHUNK_ESM)
+
+    # ——— 2) PREDIÇÕES ———
+    with st.spinner("🧠 A fazer predições…"):
+        y_pb  = mlp_pb.predict(emb_pb)
+        y_bfd = mlp_bfd.predict(emb_bfd)
+        y_esm = mlp_esm.predict(emb_esm)[:, :597]
+        X     = np.concatenate([y_pb, y_bfd, y_esm], axis=1)
+        y_ens = stacking.predict(X)
+
+    # ——— 3) MOSTRAR RESULTADOS ———
+    def mostrar(tag, y_pred):
+        with st.expander(tag, expanded=(tag == "Ensemble (Stacking)")):
             hits = mlb.inverse_transform((y_pred >= THRESH).astype(int))[0]
             st.markdown(f"**GO terms com prob ≥ {THRESH}**")
             st.code("\n".join(hits) if hits else "— nenhum —")
-
             st.markdown(f"**Top {TOP_N} GO terms mais prováveis**")
-            top_idx = np.argsort(-y_pred[0])[:TOP_N]
-            for i in top_idx:
-                st.write(f"{GO_TERMS[i]} : {y_pred[0][i]:.4f}")
-
-    # 4) OUTPUT
-    show_results("ProtBERT (MLP)",       y_pb)
-    show_results("ProtBERT-BFD (MLP)",   y_bfd)
-    show_results("ESM-2 (MLP)",          y_esm)
-    show_results("Ensemble (Stacking)",  y_ens)
+            for i in np.argsort(-y_pred[0])[:TOP_N]:
+                st.write(f"{GO[i]} : {y_pred[0][i]:.4f}")
+
+    mostrar("ProtBERT (MLP)",      y_pb)
+    mostrar("ProtBERT-BFD (MLP)",  y_bfd)
+    mostrar("ESM-2 (MLP)",         y_esm)
+    mostrar("Ensemble (Stacking)", y_ens)