Update app.py

app.py

@@ -11,7 +11,13 @@ from typing import Dict, List, Optional, Tuple
 
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoTokenizer, T5EncoderModel, T5Tokenizer
+from transformers import (
+    AutoModel,
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    T5EncoderModel,
+    T5Tokenizer,
+)
 
 APP_TITLE = "Protein Embedding"
 
@@ -252,12 +258,16 @@ class SingleModelRunner:
             from esm.models.esmc import ESMC
             self.model = ESMC.from_pretrained(spec.model_id).to(target_device)
             self.model.eval()
+            self.tokenizer = None
 
         elif spec.family == "prosst":
             ensure_prosst_repo()
 
-            self.tokenizer = AutoTokenizer.from_pretrained(spec.tokenizer_id, trust_remote_code=True)
-            self.model = AutoModel.from_pretrained(
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                spec.tokenizer_id,
+                trust_remote_code=True,
+            )
+            self.model = AutoModelForMaskedLM.from_pretrained(
                 spec.model_id,
                 trust_remote_code=True,
                 output_hidden_states=True,
@@ -289,6 +299,7 @@ def embed_hf_encoder(seq: str) -> torch.Tensor:
         truncation=False,
     )
     enc = {k: v.to(RUNNER.device) for k, v in enc.items()}
+
     out = RUNNER.model(**{k: v for k, v in enc.items() if k != "special_tokens_mask"})
     hidden = out.last_hidden_state[0]
 
@@ -312,6 +323,7 @@ def embed_t5_encoder(seq: str) -> torch.Tensor:
         truncation=False,
     )
     enc = {k: v.to(RUNNER.device) for k, v in enc.items()}
+
     out = RUNNER.model(**{k: v for k, v in enc.items() if k != "special_tokens_mask"})
     hidden = out.last_hidden_state[0]
 
@@ -354,9 +366,12 @@ def embed_esmc(seq: str) -> torch.Tensor:
     raise ValueError(f"ESMC returned shape {tuple(emb.shape)} for sequence length {len(seq)}.")
 
 
-def get_sst_tokens(seq: str):
+def get_sst_tokens(seq: str) -> List[int]:
     sst = RUNNER.sst_predictor.predict(seq)
 
+    print("SST raw type:", type(sst))
+    print("SST raw repr:", repr(sst)[:500])
+
     if isinstance(sst, str):
         tokens = [int(x) for x in sst.strip().split()]
     elif isinstance(sst, torch.Tensor):
@@ -372,7 +387,6 @@ def get_sst_tokens(seq: str):
 
     tokens = [int(x) for x in tokens]
 
-    # 尽量规整到 L
     if len(tokens) == len(seq) + 2:
         tokens = tokens[1:-1]
     elif len(tokens) == len(seq) + 1:
@@ -383,6 +397,9 @@ def get_sst_tokens(seq: str):
     if len(tokens) != len(seq):
         raise ValueError(f"SST token length mismatch: got {len(tokens)}, expected {len(seq)}")
 
+    print("SST final length:", len(tokens))
+    print("SST first 30:", tokens[:30])
+
     return tokens
 
 
@@ -400,8 +417,6 @@ def embed_prosst(seq: str) -> Tuple[torch.Tensor, List[int]]:
     )
     seq_enc = {k: v.to(RUNNER.device) for k, v in seq_enc.items()}
 
-    # ProSST 常见做法是把结构 token 当作额外输入 ids
-    # 这里直接构建 [1, L] LongTensor
     sst_ids = torch.tensor([sst_tokens], dtype=torch.long, device=RUNNER.device)
 
     tried = []
@@ -411,9 +426,15 @@ def embed_prosst(seq: str) -> Tuple[torch.Tensor, List[int]]:
                 input_ids=seq_enc["input_ids"],
                 attention_mask=seq_enc.get("attention_mask", None),
                 output_hidden_states=True,
+                return_dict=True,
                 **{kw: sst_ids},
             )
+
+            if getattr(out, "hidden_states", None) is None:
+                raise RuntimeError("ProSST output has no hidden_states")
+
             hidden = out.hidden_states[-1][0]
+
             emb = normalize_to_Ld(
                 hidden=hidden,
                 expected_len=len(seq),
@@ -421,10 +442,13 @@ def embed_prosst(seq: str) -> Tuple[torch.Tensor, List[int]]:
                 attention_mask=seq_enc.get("attention_mask", None)[0] if seq_enc.get("attention_mask", None) is not None else None,
             )
             return emb.detach().cpu().float(), sst_tokens
+
         except Exception as e:
             tried.append(f"{kw}: {repr(e)}")
 
-    raise RuntimeError("Failed to run ProSST with known structure-token arg names: " + " | ".join(tried))
+    raise RuntimeError(
+        "Failed to run ProSST with known structure-token arg names: " + " | ".join(tried)
+    )
 
 
 def embed_one_sequence(seq: str):