Update app.py

app.py

@@ -1,18 +1,16 @@
-import os
-import re
-import numpy as np
-import torch
-import joblib
-import streamlit as st
+# -------------------------------------------------------------------------------------------------
+#  app.py  –  Streamlit app para predição de GO:MF
+#  Versão: usa ProtBERT & ProtBERT-BFD fine-tuned (melvinalves/FineTune) + ESM-2 base
+# -------------------------------------------------------------------------------------------------
+import os, re, numpy as np, torch, joblib, streamlit as st
+from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModel
-from huggingface_hub import hf_hub_download
 from keras.models import load_model
 from goatools.obo_parser import GODag
 
-
-from huggingface_hub import login
-
+# ———————————————————  AUTHENTICAÇÃO  ——————————————————— #
 login(os.environ["HF_TOKEN"])
+
 # ———————————————————  CONFIG  ——————————————————— #
 SPACE_ID   = "melvinalves/protein_function_prediction"
 TOP_N      = 10
@@ -20,36 +18,63 @@ THRESH     = 0.37
 CHUNK_PB   = 512
 CHUNK_ESM  = 1024
 
+# Repositórios dos modelos
+FINETUNED_PB   = ("melvinalves/FineTune", "fineTunedProtbert")
+FINETUNED_BFD  = ("melvinalves/FineTune", "fineTunedProtbertbfd")
+BASE_ESM       = "facebook/esm2_t33_650M_UR50D"
+
 # ———————————————————  HELPERS  ——————————————————— #
 @st.cache_resource
 def download_file(path):
+    """Ficheiros pequenos guardados no repositório do Space (≤1 GB total)."""
+    from huggingface_hub import hf_hub_download
     return hf_hub_download(repo_id=SPACE_ID, repo_type="space", filename=path)
 
 @st.cache_resource
 def load_keras(name):
+    """Carrega modelos Keras (MLPs + stacking)."""
     return load_model(download_file(f"models/{name}"), compile=False)
 
 @st.cache_resource
-def load_hf_encoder(model):
-    tok = AutoTokenizer.from_pretrained(model, do_lower_case=False)
-    mdl = AutoModel.from_pretrained(model)
+def load_hf_encoder(repo_id, subfolder=None, base_tok="Rostlab/prot_bert"):
+    """
+    Carrega um encoder HF (PyTorch) – se existir apenas tf_model.h5 no repo,
+    usa from_tf=True para converter on-the-fly.
+    """
+    tok = AutoTokenizer.from_pretrained(base_tok, do_lower_case=False)
+    mdl = AutoModel.from_pretrained(
+        repo_id,
+        subfolder=subfolder,
+        from_tf=True,   # converte pesos TF se necessário
+    )
     mdl.eval()
     return tok, mdl
 
-def embed_seq(model, seq, chunk):
-    tok, mdl = load_hf_encoder(model)
+def embed_seq(model_ref, seq, chunk):
+    """
+    Extrai embedding médio (CLS) para sequências grandes usando chunks.
+    - model_ref  pode ser string (modelo base) ou tuple (repo_id, subfolder) p/ fine-tuned.
+    """
+    if isinstance(model_ref, tuple):
+        tok, mdl = load_hf_encoder(*model_ref)
+    else:
+        # mantém o tokenizer apropriado
+        base_tok = "Rostlab/prot_bert" if "prot_bert" in model_ref else model_ref
+        tok, mdl = load_hf_encoder(model_ref, base_tok=base_tok)
+
     parts = [seq[i:i+chunk] for i in range(0, len(seq), chunk)]
-    vecs  = []
+    vecs = []
     for p in parts:
+        tokens = tok(" ".join(p), return_tensors="pt", truncation=False)
         with torch.no_grad():
-            out = mdl(**tok(" ".join(p), return_tensors="pt", truncation=False))
-        vecs.append(out.last_hidden_state[:, 0, :].squeeze().numpy())
+            out = mdl(**{k: v.to(mdl.device) for k, v in tokens.items()})
+        vecs.append(out.last_hidden_state[:, 0, :].cpu().numpy())
     return np.mean(vecs, axis=0, keepdims=True)
 
 @st.cache_resource
 def load_go_info():
     obo_path = download_file("data/go.obo")
-    dag = GODag(obo_path, optional_attrs=['defn'])
+    dag = GODag(obo_path, optional_attrs=["defn"])
     return {tid: (term.name, term.defn) for tid, term in dag.items()}
 
 GO_INFO = load_go_info()
@@ -67,9 +92,7 @@ GO       = mlb.classes_
 st.title("Predição de Funções Moleculares de Proteínas")
 
 st.markdown(
-    """
-    <style> textarea { font-size: 0.9rem !important; } </style>
-    """,
+    "<style> textarea { font-size: 0.9rem !important; } </style>",
     unsafe_allow_html=True,
 )
 
@@ -78,58 +101,55 @@ predict_clicked = st.button("Prever GO terms")
 
 # ———————————————————  PARSE DE MÚLTIPLAS SEQUÊNCIAS  ——————————————————— #
 def parse_fasta_multiple(fasta_str):
-    entries = fasta_str.strip().split(">")
-    parsed = []
-
+    entries, parsed = fasta_str.strip().split(">"), []
     for i, entry in enumerate(entries):
         if not entry.strip():
             continue
-
         lines = entry.strip().splitlines()
-
-        # Verifica se estamos num bloco com '>'
         if i > 0:
             header = lines[0].strip()
-            seq = "".join(l.strip() for l in lines[1:]).replace(" ", "").upper()
+            seq = "".join(lines[1:]).replace(" ", "").upper()
         else:
-            # Entrada sem '>', trata tudo como sequência
             header = f"Seq_{i+1}"
-            seq = "".join(l.strip() for l in lines).replace(" ", "").upper()
-
+            seq = "".join(lines).replace(" ", "").upper()
         if seq:
             parsed.append((header, seq))
     return parsed
 
+# ———————————————————  INFERÊNCIA  ——————————————————— #
 if predict_clicked:
     parsed_seqs = parse_fasta_multiple(fasta_input)
-
     if not parsed_seqs:
         st.warning("Não foi possível encontrar nenhuma sequência válida.")
         st.stop()
 
     for header, seq in parsed_seqs:
         with st.spinner(f"A processar {header}… (pode demorar alguns minutos)"):
-            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)
+            # ——— Embeddings ——— #
+            emb_pb  = embed_seq(FINETUNED_PB,  seq, CHUNK_PB)
+            emb_bfd = embed_seq(FINETUNED_BFD, seq, CHUNK_PB)
+            emb_esm = embed_seq(BASE_ESM,       seq, CHUNK_ESM)
 
+            # ——— Predições dos MLPs ——— #
             y_pb  = mlp_pb.predict(emb_pb)
             y_bfd = mlp_bfd.predict(emb_bfd)
             y_esm = mlp_esm.predict(emb_esm)[:, :597]
+
+            # ——— Stacking ——— #
             X     = np.concatenate([y_pb, y_bfd, y_esm], axis=1)
             y_ens = stacking.predict(X)
 
         # ———————————————————  RESULTADOS ——————————————————— #
-        def mostrar(tag, y_pred):
+        def mostrar_resultados(tag, y_pred):
             with st.expander(tag, expanded=True):
                 hits = mlb.inverse_transform((y_pred >= THRESH).astype(int))[0]
                 st.markdown(f"**GO terms com prob ≥ {THRESH}**")
                 if hits:
                     for go_id in hits:
                         name, defin = GO_INFO.get(go_id, ("— sem nome —", ""))
-                        limpa_def = re.sub(r'^\s*"?(.+?)"?\s*(\[[^\]]*\])?\s*$', r'\1', defin or "")
+                        limp = re.sub(r'^\s*"?(.+?)"?\s*(\[[^\]]*\])?\s*$', r'\1', defin or "")
                         st.write(f"**{go_id} — {name}**")
-                        st.caption(limpa_def)
+                        st.caption(limp)
                 else:
                     st.code("— nenhum —")
 
@@ -139,9 +159,8 @@ if predict_clicked:
                     name, _ = GO_INFO.get(go_id, ("", ""))
                     st.write(f"{go_id} — {name} : {y_pred[0][idx]:.4f}")
 
-        # ———————————————————  ESCOLHE QUAIS MOSTRAR ——————————————————— #
-
-        # mostrar(f"{header} — ProtBERT (MLP)",      y_pb)
-        # mostrar(f"{header} — ProtBERT-BFD (MLP)",  y_bfd)
-        # mostrar(f"{header} — ESM-2 (MLP)",         y_esm)
-        mostrar(f"{header}", y_ens)
+        # Mostrar apenas ensemble (descomenta se quiseres os individuais)
+        # mostrar_resultados(f"{header} — ProtBERT",      y_pb)
+        # mostrar_resultados(f"{header} — ProtBERT-BFD",  y_bfd)
+        # mostrar_resultados(f"{header} — ESM-2",         y_esm)
+        mostrar_resultados(header, y_ens)