Update app.py

app.py

@@ -24,6 +24,20 @@ def load_hf_model(name):
     model.eval()
     return tokenizer, model
 
+# ---------- Cache dos modelos locais ----------
+@st.cache_resource
+def load_local_model(path):
+    return load_model(path, compile=False)
+
+mlp_pb  = load_local_model(os.path.join(MODELS_DIR, "mlp_protbert.keras"))
+mlp_bfd = load_local_model(os.path.join(MODELS_DIR, "mlp_protbertbfd.keras"))
+mlp_esm = load_local_model(os.path.join(MODELS_DIR, "mlp_esm2.keras"))
+stacking = load_local_model(os.path.join(MODELS_DIR, "modelo_ensemble_stack.keras"))
+
+# ---------- Carregar MultiLabelBinarizer ----------
+mlb = joblib.load(MLB_PATH)
+go_terms = mlb.classes_
+
 # ---------- Função para gerar embedding por chunk ----------
 def embed_sequence(model_name, seq, chunk_size):
     tokenizer, model = load_hf_model(model_name)
@@ -44,16 +58,6 @@ def embed_sequence(model_name, seq, chunk_size):
 
     return np.mean(embeddings, axis=0, keepdims=True)
 
-# ---------- Carregar modelos ----------
-mlp_pb  = load_model(os.path.join(MODELS_DIR, "mlp_protbert.keras"), compile=False)
-mlp_bfd = load_model(os.path.join(MODELS_DIR, "mlp_protbertbfd.keras"), compile=False)
-mlp_esm = load_model(os.path.join(MODELS_DIR, "mlp_esm2.keras"), compile=False)
-stacking = load_model(os.path.join(MODELS_DIR, "modelo_ensemble_stack.keras"), compile=False)
-
-# ---------- Carregar MultiLabelBinarizer ----------
-mlb = joblib.load(MLB_PATH)
-go_terms = mlb.classes_
-
 # ---------- Interface Streamlit ----------
 st.title("Predição de Funções de Proteínas")
 
@@ -68,17 +72,17 @@ if st.button("Prever GO terms"):
     if not seq:
         st.warning("Por favor, insere uma sequência válida.")
     else:
-        st.write("A gerar embeddings por chunks...")
+        st.write("🔄 A gerar embeddings...")
 
         emb_pb  = embed_sequence("Rostlab/prot_bert", seq, CHUNK_PB)
         emb_bfd = embed_sequence("Rostlab/prot_bert_bfd", seq, CHUNK_PB)
         emb_esm = embed_sequence("facebook/esm2_t33_650M_UR50D", seq, CHUNK_ESM)
 
-        st.write("A fazer predições base...")
+        st.write("🧠 A fazer predições com cada modelo...")
 
-        y_pb  = mlp_pb.predict(emb_pb)[:, :597]
-        y_bfd = mlp_bfd.predict(emb_bfd)[:, :597]
-        y_esm = mlp_esm.predict(emb_esm)[:, :597]
+        y_pb  = mlp_pb.predict(emb_pb)
+        y_bfd = mlp_bfd.predict(emb_bfd)
+        y_esm = mlp_esm.predict(emb_esm)
 
         X_stack = np.concatenate([y_pb, y_bfd, y_esm], axis=1)
         y_pred = stacking.predict(X_stack)
@@ -94,4 +98,3 @@ if st.button("Prever GO terms"):
         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}")
-