Update app.py

app.py

@@ -211,12 +211,11 @@ def embed_sequence(seq: str) -> np.ndarray:
             out = model(**enc, output_hidden_states=True)
         hs = out.hidden_states
         mask = enc["attention_mask"].unsqueeze(-1).float()
-        mvecs = []
-        for li in [8, 10, 11]:
-            h = hs[li]
-            mv = (h * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
-            mvecs.append(mv.squeeze(0).cpu().numpy())
-        return np.concatenate(mvecs)
+        
+        # Grab the FINAL layer (-1) instead of hardcoding [8, 10, 11]
+        h = hs[-1]
+        mv = (h * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
+        return mv.squeeze(0).cpu().numpy()
 
     seq = seq.strip()
     if len(seq) <= MAX: