Update app.py

app.py

@@ -1,93 +1,102 @@
-import os, pathlib
+import os
 import numpy as np
-import torch, joblib, streamlit as st
+import torch
+import joblib
+import streamlit as st
 from transformers import AutoTokenizer, AutoModel
 from huggingface_hub import hf_hub_download
-from tensorflow.keras.models import load_model
+from keras.models import load_model
 
 # ---------- Configuração ----------
-SPACE_ID   = "melvinalves/protein_function_prediction"   # id deste Space
-TOP_N      = 10
-CHUNK_PB   = 512
-CHUNK_ESM  = 1024
+SPACE_ID = "melvinalves/protein_function_prediction"
+TOP_N = 10
+CHUNK_PB = 512
+CHUNK_ESM = 1024
 
-# ---------- Helpers ----------
+# ---------- Cache de downloads ----------
 @st.cache_resource
-def hf_cached(fname: str):
-    """Faz download de um ficheiro (LFS ou não) e devolve o caminho local."""
-    return hf_hub_download(
-        repo_id   = SPACE_ID,
-        repo_type = "space",
-        filename  = fname,        # ex: "models/mlp_protbert.keras"
+def download_model_file(filename):
+    local_path = hf_hub_download(
+        repo_id=SPACE_ID,
+        repo_type="space",
+        filename=f"models/{filename}",
     )
+    print(f"📦 {filename} → {os.path.getsize(local_path)} bytes")
+    return local_path
 
 @st.cache_resource
-def load_hf_model(model_name: str):
-    tok  = AutoTokenizer.from_pretrained(model_name, do_lower_case=False)
-    mdl  = AutoModel.from_pretrained(model_name)
-    mdl.eval()
-    return tok, mdl
+def load_hf_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=False)
+    model = AutoModel.from_pretrained(model_name)
+    model.eval()
+    return tokenizer, model
 
 @st.cache_resource
-def load_local_model(fname: str):
-    full_path = hf_cached(f"models/{fname}")
-    print(">>", pathlib.Path(full_path).name,
-      "tamanho:", os.path.getsize(full_path), "bytes")
-    return load_model(full_path, compile=False)
-
-# ---------- Carregamento dos modelos locais (.keras) ----------
-mlp_pb   = load_local_model("mlp_protbert.keras")
-mlp_bfd  = load_local_model("mlp_protbertbfd.keras")
-mlp_esm  = load_local_model("mlp_esm2.keras")          # 602 saídas – cortamos depois
-stacking = load_local_model("ensemble_stacking.keras") # espera 1791 entradas
-
-# ---------- Label binarizer ----------
-mlb      = joblib.load(hf_cached("data/mlb_597.pkl"))
+def load_keras_model(filename):
+    path = download_model_file(filename)
+    return load_model(path, compile=False)
+
+# ---------- Carregar modelos ----------
+mlp_pb   = load_keras_model("mlp_protbert.keras")
+mlp_bfd  = load_keras_model("mlp_protbertbfd.keras")
+mlp_esm  = load_keras_model("mlp_esm2.keras")
+stacking = load_keras_model("ensemble_stacking.keras")
+
+# ---------- Carregar MultiLabelBinarizer ----------
+mlb = joblib.load(hf_hub_download(
+    repo_id=SPACE_ID,
+    repo_type="space",
+    filename="data/mlb_597.pkl"
+))
 go_terms = mlb.classes_
 
-# ---------- Função de embedding ----------
-def embed_sequence(model_name: str, seq: str, chunk: int) -> np.ndarray:
-    tok, mdl = load_hf_model(model_name)
-    fmt = lambda s: " ".join(list(s))
-    parts = [seq[i:i+chunk] for i in range(0, len(seq), chunk)]
-    vecs  = []
-    for p in parts:
+# ---------- Função para gerar embeddings ----------
+def embed_sequence(model_name, seq, chunk_size):
+    tokenizer, model = load_hf_model(model_name)
+    def format_seq(s): return " ".join(list(s))
+    chunks = [seq[i:i+chunk_size] for i in range(0, len(seq), chunk_size)]
+    embeddings = []
+    for chunk in chunks:
+        formatted = format_seq(chunk)
+        inputs = tokenizer(formatted, return_tensors="pt", truncation=True)
         with torch.no_grad():
-            out = mdl(**tok(fmt(p), return_tensors="pt", truncation=True))
-        vecs.append(out.last_hidden_state[:, 0, :].squeeze().numpy())
-    return np.mean(vecs, axis=0, keepdims=True)
+            outputs = model(**inputs)
+        cls_embedding = outputs.last_hidden_state[:, 0, :].squeeze().numpy()
+        embeddings.append(cls_embedding)
+    return np.mean(embeddings, axis=0, keepdims=True)
 
-# ---------- Interface ----------
+# ---------- Interface Streamlit ----------
 st.title("🔬 Predição de Funções de Proteínas")
 
-src = st.text_area("Insere a sequência FASTA:", height=200)
+user_input = st.text_area("Insere a sequência FASTA:", height=200)
 
-if src and st.button("Prever GO terms"):
-    # limpar FASTA
-    seq = "\n".join(l for l in src.splitlines() if not l.startswith(">"))
-    seq = seq.replace(" ", "").replace("\n", "").upper()
-    if not seq:
-        st.warning("Sequência vazia.")
+if user_input and st.button("Prever GO terms"):
+    # Limpar sequência FASTA
+    sequence = "\n".join([line for line in user_input.splitlines() if not line.startswith(">")])
+    sequence = sequence.replace(" ", "").replace("\n", "").strip().upper()
+    if not sequence:
+        st.warning("Por favor, insere uma sequência válida.")
         st.stop()
 
     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)
+    emb_pb  = embed_sequence("Rostlab/prot_bert",            sequence, CHUNK_PB)
+    emb_bfd = embed_sequence("Rostlab/prot_bert_bfd",        sequence, CHUNK_PB)
+    emb_esm = embed_sequence("facebook/esm2_t33_650M_UR50D", sequence, CHUNK_ESM)
 
-    st.write("🧠 A prever com cada modelo…")
+    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 para alinhar
+    y_esm = mlp_esm.predict(emb_esm)[:, :597]  # garantir alinhamento
 
-    X_stack = np.concatenate([y_pb, y_bfd, y_esm], axis=1)   # (1, 1791)
-    y_pred  = stacking.predict(X_stack)
+    X_stack = np.concatenate([y_pb, y_bfd, y_esm], axis=1)
+    y_pred = stacking.predict(X_stack)
 
-    # ---------- Resultados ----------
-    st.subheader("GO terms com probabilidade ≥ 0.5")
+    st.subheader("🎯 GO terms com probabilidade ≥ 0.5")
     hits = mlb.inverse_transform((y_pred >= 0.5).astype(int))[0]
-    st.code("\n".join(hits) or "— nenhum —")
+    st.code("\n".join(hits) if hits else "— nenhum —")
+
+    st.subheader(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}")
 
-    st.subheader(f"Top {TOP_N} GO terms mais prováveis")
-    for idx in np.argsort(-y_pred[0])[:TOP_N]:
-        st.write(f"{go_terms[idx]} : {y_pred[0][idx]:.4f}")