Update app.py

app.py

@@ -9,8 +9,10 @@ from transformers import (
     LogitsProcessorList,
 )
 
+# ── Config ──────────────────────────────────────────────────────────────────────
 MODEL_ID = os.getenv("MODEL_ID", "llm-rna-api-rmit/rna-structure-model")
 
+# ── Model loading ───────────────────────────────────────────────────────────────
 @lru_cache(maxsize=1)
 def _load_model_and_tokenizer():
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -21,58 +23,63 @@ def _load_model_and_tokenizer():
         device_map="auto" if device == "cuda" else None,
     )
     model.eval()
+    if device != "cuda":
+        model.to(device)
     return tokenizer, model, device
 
-# --- Utility helpers ---
+# ── Utility helpers ─────────────────────────────────────────────────────────────
 def _char_token_id(tokenizer, ch: str) -> int:
+    # Prefer an exact single-char token if it exists
     ids = tokenizer.encode(ch, add_special_tokens=False)
     for tid in ids:
         if tokenizer.decode([tid]) == ch:
             return tid
+    # Fallback: scan vocab for an exact decode match
     for tid in range(len(tokenizer)):
         if tokenizer.decode([tid]) == ch:
             return tid
-    raise ValueError(f"Could not find token id for {ch}")
+    raise ValueError(f"Could not find token id for {ch!r}")
 
 def _can_pair(a, b, allow_gu=True):
-    if (a,b) in [("A","U"),("U","A"),("G","C"),("C","G")]:
+    if (a, b) in [("A","U"),("U","A"),("G","C"),("C","G")]:
         return True
-    if allow_gu and (a,b) in [("G","U"),("U","G")]:
+    if allow_gu and (a, b) in [("G","U"),("U","G")]:
         return True
     return False
 
 def _precompute_can_open(seq, min_loop=3, allow_gu=True):
-    n=len(seq)
-    can=[False]*n
+    n = len(seq)
+    can = [False] * n
     for i in range(n):
-        for j in range(i+min_loop+1,n):
-            if _can_pair(seq[i],seq[j],allow_gu):
-                can[i]=True
+        for j in range(i + min_loop + 1, n):
+            if _can_pair(seq[i], seq[j], allow_gu):
+                can[i] = True
                 break
     return can
 
-# --- constrained processor ---
-# --- constrained processor ---
+# ── Constrained processor ───────────────────────────────────────────────────────
 class BalancedParenProcessor(LogitsProcessor):
+    """
+    Restricts next token to one of: '(', ')' or '.', while maintaining balance
+    and leaving room to close opened stems. No dot bias by default.
+    """
     def __init__(self, lp_id, rp_id, dot_id, total_len, can_open,
                  dot_bias=0.0, paren_penalty=0.0, window=5):
         self.lp_id, self.rp_id, self.dot_id = lp_id, rp_id, dot_id
         self.total_len = total_len
         self.step = 0
         self.depth = 0
-        self.history=[]
+        self.history = []
         self.can_open = can_open
-        self.dot_bias=dot_bias
-        self.paren_penalty=paren_penalty
-        self.window=window
+        self.dot_bias = dot_bias
+        self.paren_penalty = paren_penalty
+        self.window = window
 
     def __call__(self, input_ids, scores):
-        # restrict to only three tokens
+        # Mask everything except allowed tokens
         mask = torch.full_like(scores, float("-inf"))
         remaining = self.total_len - self.step
         allowed = []
-
-        # If we must close to avoid running out of room, force )
         must_close = (remaining == self.depth and self.depth > 0)
         pos = self.step
 
@@ -81,22 +88,18 @@ class BalancedParenProcessor(LogitsProcessor):
         else:
             if self.depth > 0:
                 allowed.append(self.rp_id)
-
-            # allow opening if there will still be room to close later
-            # (be a bit less strict than remaining-1 > depth to encourage stems)
+            # Allow opening if there will still be room to close later
             if remaining - 2 >= self.depth and pos < len(self.can_open) and self.can_open[pos]:
                 allowed.append(self.lp_id)
-
             allowed.append(self.dot_id)
 
         mask[:, allowed] = 0.0
         scores = scores + mask
 
-        # (no dot boost by default)
         if self.dot_bias != 0.0:
             scores[:, self.dot_id] += self.dot_bias
 
-        # optional mild anti-run regularizer
+        # Optional mild anti-run for long paren streaks
         if self.paren_penalty and len(self.history) >= self.window and all(
             t in (self.lp_id, self.rp_id) for t in self.history[-self.window:]
         ):
@@ -105,19 +108,49 @@ class BalancedParenProcessor(LogitsProcessor):
 
         return scores
 
-# --- generator ---
+    def update(self, tok):
+        if tok == self.lp_id:
+            self.depth += 1
+        elif tok == self.rp_id:
+            self.depth = max(0, self.depth - 1)
+        self.history.append(tok)
+        self.step += 1
+
+def _top_p_sample(logits, top_p=0.9, temperature=0.8):
+    logits = logits / temperature
+    probs = torch.softmax(logits, dim=-1)
+    sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+    cumsum = torch.cumsum(sorted_probs, dim=-1)
+    mask = cumsum > top_p
+    mask[..., 0] = False
+    sorted_probs[mask] = 0
+    sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+    idx = torch.multinomial(sorted_probs, 1)
+    return sorted_idx.gather(-1, idx).squeeze(-1)
+
+# ── Generator ───────────────────────────────────────────────────────────────────
 def _generate_db(seq):
     tok, model, device = _load_model_and_tokenizer()
     n = len(seq)
     prompt = f"RNA: {seq}\nDot-bracket (exactly {n} characters using only '(' ')' '.'):\n"
-    lp = _char_token_id(tok, "("); rp = _char_token_id(tok, ")"); dot = _char_token_id(tok, ".")
-    can = _precompute_can_open(seq, min_loop=3)  # try 2 if you still get few stems
+
+    lp = _char_token_id(tok, "(")
+    rp = _char_token_id(tok, ")")
+    dot = _char_token_id(tok, ".")
+
+    # Helpful to verify once in logs
+    print("Token IDs:", {"(": lp, ")": rp, ".": dot})
+
+    can = _precompute_can_open(seq, min_loop=3, allow_gu=True)
     proc = BalancedParenProcessor(lp, rp, dot, n, can, dot_bias=0.0, paren_penalty=0.0)
     procs = LogitsProcessorList([proc])
-    inputs = tok(prompt, return_tensors="pt").to(device)
+
+    inputs = tok(prompt, return_tensors="pt")
+    # Keep tensors on the same device as the model
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
     cur = inputs["input_ids"]
-    generated = []
 
+    generated = []
     with torch.no_grad():
         for _ in range(n):
             out = model(cur)
@@ -125,11 +158,11 @@ def _generate_db(seq):
             for p in procs:
                 logits = p(cur, logits)
             next_id = _top_p_sample(logits, top_p=0.9, temperature=0.8)
-            next_id = next_id.to(device)
             tokid = next_id.item()
             generated.append(tokid)
             proc.update(tokid)
-            cur = torch.cat([cur, next_id.view(1, 1)], dim=1)
+            # Make sure we append on the SAME device as cur/model
+            cur = torch.cat([cur, next_id.view(1, 1).to(cur.device)], dim=1)
 
     text = tok.decode(generated, skip_special_tokens=True)
     db = "".join(c for c in text if c in "().")[:n]
@@ -137,45 +170,56 @@ def _generate_db(seq):
         db = (db + "." * n)[:n]
     return db
 
-
-# --- structural element translation ---
+# ── Structural element translation ──────────────────────────────────────────────
 def dotbracket_to_structural(dot_str):
-    if not dot_str: return "<start><external_loop><end>"
-    res=["<start>"];depth=0;i=0;n=len(dot_str)
-    def add(tag): 
-        if res[-1]!=tag:res.append(tag)
-    while i<n:
-        c=dot_str[i]
-        if c==".":
-            j=i
-            while j<n and dot_str[j]==".":
-                j+=1
-            nextc=dot_str[j] if j<n else None
-            tag="<external_loop>" if depth==0 else ("<hairpin>" if nextc==")" else "<internal_loop>")
-            add(tag);i=j;continue
-        if c=="(":
-            add("<stem>");depth+=1
-        else:
-            add("<stem>");depth=max(0,depth-1)
-        i+=1
+    if not dot_str or not isinstance(dot_str, str):
+        return "<start><external_loop><end>"
+    res = ["<start>"]; depth = 0; i = 0; n = len(dot_str)
+
+    def add(tag):
+        if res[-1] != tag:
+            res.append(tag)
+
+    while i < n:
+        c = dot_str[i]
+        if c == ".":
+            j = i
+            while j < n and dot_str[j] == ".":
+                j += 1
+            nextc = dot_str[j] if j < n else None
+            tag = "<external_loop>" if depth == 0 else ("<hairpin>" if nextc == ")" else "<internal_loop>")
+            add(tag); i = j; continue
+        if c == "(":
+            add("<stem>"); depth += 1
+        else:  # ')'
+            add("<stem>"); depth = max(0, depth - 1)
+        i += 1
+
     res.append("<end>")
     return "".join(res)
 
-# --- Gradio wrapper ---
-def predict(seq):
-    seq=(seq or "").strip().upper()
-    if not seq or not set(seq)<={"A","U","C","G"}:
-        return "Please enter an RNA sequence (A/U/C/G)."
-    db=_generate_db(seq)
-    return dotbracket_to_structural(db)
+# ── Gradio wrapper ──────────────────────────────────────────────────────────────
+import traceback
 
-demo=gr.Interface(
+def predict(seq):
+    try:
+        seq = (seq or "").strip().upper()
+        if not seq or not set(seq) <= {"A", "U", "C", "G"}:
+            return "Please enter an RNA sequence (A/U/C/G)."
+        db = _generate_db(seq)
+        return dotbracket_to_structural(db)
+    except Exception as e:
+        # Print full traceback to Space logs and show a concise error in UI
+        traceback.print_exc()
+        return f"Error: {type(e).__name__}: {e}"
+
+demo = gr.Interface(
     fn=predict,
-    inputs=gr.Textbox(lines=4,label="RNA Sequence (A/U/C/G)",value="GGGAAUCC"),
-    outputs=gr.Textbox(lines=6,label="Predicted Structural Elements"),
+    inputs=gr.Textbox(lines=4, label="RNA Sequence (A/U/C/G)", value="GGGAAUCC"),
+    outputs=gr.Textbox(lines=6, label="Predicted Structural Elements"),
     title="RNA Structure Predictor",
     description="Outputs <start>, <stem>, <hairpin>, <internal_loop>, <external_loop>, <end>."
 )
 
-if __name__=="__main__":
+if __name__ == "__main__":
     demo.launch()