Update app.py

app.py

@@ -3,7 +3,12 @@ import re
 import torch
 import gradio as gr
 from functools import lru_cache
-from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    LogitsProcessor,
+    LogitsProcessorList,
+)
 
 MODEL_ID = os.getenv("MODEL_ID", "llm-rna-api-rmit/rna-structure-model")
 
@@ -19,70 +24,123 @@ def _load_model_and_tokenizer():
     model.eval()
     return tokenizer, model, device
 
+# ---------- Prompt (few-shot to reduce "all dots") ----------
+FEWSHOT = """Return ONLY the dot-bracket structure as one line with the same length as RNA.
+
+RNA: GCGCGAAAACGCGC
+Dot-bracket: (((((....)))))
+
+RNA: GGGAAAUCCCU
+Dot-bracket: (((...)))
+"""
+
 def _make_prompt(seq: str) -> str:
     n = len(seq)
     return (
-        f"RNA: {seq}\n"
-        f"Output ONLY the RNA secondary structure in dot-bracket notation, exactly {n} characters long, "
-        f"using only '(' ')' and '.'.\n"
-        f"Dot-bracket:"
+        FEWSHOT
+        + f"\nRNA: {seq}\n"
+        + f"Dot-bracket (exactly {n} characters using only '(' ')' '.'):"
     )
 
-def _generate(prompt: str, max_new_tokens: int = 256):
+# ---------- Robust char→token id ----------
+def _char_token_id(tokenizer, ch: str) -> int:
+    # Try simple path
+    ids = tokenizer.encode(ch, add_special_tokens=False)
+    if ids:
+        # prefer single-token mapping when it decodes back to the same char
+        for tid in ids:
+            if tokenizer.decode([tid]) == ch:
+                return tid
+        return ids[-1]
+    # Fallback: scan vocab for a token that decodes to ch
+    for tid in range(len(tokenizer)):
+        if tokenizer.decode([tid]) == ch:
+            return tid
+    raise ValueError(f"Could not find token id for char {ch!r}")
+
+# ---------- Constrained + biased generation ----------
+class AllowOnlyAndBias(LogitsProcessor):
+    def __init__(self, allowed_ids, bias_map):
+        self.allowed = torch.tensor(allowed_ids, dtype=torch.long)
+        self.bias_map = {int(k): float(v) for k, v in bias_map.items()}
+    def __call__(self, input_ids, scores):
+        # mask everything else
+        scores[:] = float("-inf")
+        scores[:, self.allowed] = 0.0
+        # add biases to steer away from '.' and toward parentheses
+        for tid, bias in self.bias_map.items():
+            scores[:, tid] += bias
+        return scores
+
+def _generate_db(seq: str) -> str:
     tokenizer, model, device = _load_model_and_tokenizer()
+    n = len(seq)
+    prompt = _make_prompt(seq)
 
-    class StopOnNewline(StoppingCriteria):
-        def __init__(self, newline_id: int): self.newline_id = newline_id
-        def __call__(self, input_ids, scores, **kwargs):
-            return input_ids[0, -1].item() == self.newline_id
+    # get robust IDs for '(', ')', '.'
+    lp_id = _char_token_id(tokenizer, "(")
+    rp_id = _char_token_id(tokenizer, ")")
+    dot_id = _char_token_id(tokenizer, ".")
+
+    processors = LogitsProcessorList([
+        # encourage parentheses, discourage all-dots
+        AllowOnlyAndBias(
+            allowed_ids=[lp_id, rp_id, dot_id],
+            bias_map={lp_id: +1.2, rp_id: +1.2, dot_id: -0.8},
+        )
+    ])
 
-    nl_id = tokenizer.encode("\n", add_special_tokens=False)[0]
     with torch.inference_mode():
         inputs = tokenizer(prompt, return_tensors="pt").to(device)
-        outputs = model.generate(
+        out = model.generate(
             **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.0,
+            max_new_tokens=n,
+            min_new_tokens=n,
             do_sample=False,
+            temperature=0.0,
+            logits_processor=processors,
             pad_token_id=tokenizer.eos_token_id,
             eos_token_id=tokenizer.eos_token_id,
-            stopping_criteria=StoppingCriteriaList([StopOnNewline(nl_id)]),
         )
-    gen_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
-    return tokenizer.decode(gen_tokens, skip_special_tokens=True)
-
-def _extract_dotbracket_strict(text: str, length: int):
-    text = text.strip()
-    candidates = []
-
-    # Filter all dot-bracket-like substrings
-    for line in text.splitlines():
-        line = line.strip()
-        cand = "".join(c for c in line if c in "().")
-        if cand:
-            candidates.append(cand)
-
-    # Choose first one with exact or closest match
-    for cand in candidates:
-        if len(cand) == length:
-            return cand
-    if candidates:
-        # fallback: pick longest valid segment if none matches perfectly
-        return max(candidates, key=len)
-    return None
 
+    gen = out[0][inputs["input_ids"].shape[-1]:]
+    text = tokenizer.decode(gen, skip_special_tokens=True)
+    db = "".join(c for c in text if c in "().")[:n]
+    if len(db) != n:
+        db = (db + "." * n)[:n]
+    return db
+
+# ---------- Simple heuristic fallback if model gives all dots ----------
+def _is_complement(a, b):
+    return (a == "G" and b == "C") or (a == "C" and b == "G") or (a == "A" and b == "U") or (a == "U" and b == "A")
+
+def naive_hairpin(seq: str, min_loop: int = 3) -> str:
+    n = len(seq)
+    db = ["." for _ in range(n)]
+    i, j = 0, n - 1
+    while i < j - min_loop:
+        if _is_complement(seq[i], seq[j]):
+            db[i], db[j] = "(", ")"
+            i += 1
+            j -= 1
+        else:
+            # move the weaker side inward to try to find a match
+            if seq[i] in "AU" and seq[j] in "GC":
+                i += 1
+            else:
+                j -= 1
+    return "".join(db)
+
+# ---------- Dot-bracket → structural ----------
 def dotbracket_to_structural(dot_str: str) -> str:
     if not dot_str:
         return "<start><external_loop><end>"
-
-    n = len(dot_str)
     res = ["<start>"]
     depth = 0
-    i = 0
+    i, n = 0, len(dot_str)
 
-    def append_once(tag: str):
-        if res[-1] != tag:
-            res.append(tag)
+    def append_once(tag): 
+        if res[-1] != tag: res.append(tag)
 
     while i < n:
         c = dot_str[i]
@@ -91,17 +149,14 @@ def dotbracket_to_structural(dot_str: str) -> str:
             while j < n and dot_str[j] == '.':
                 j += 1
             next_char = dot_str[j] if j < n else None
-            if depth == 0:
-                label = "<external_loop>"
-            else:
-                label = "<hairpin>" if next_char == ')' else "<internal_loop>"
+            label = "<external_loop>" if depth == 0 else ("<hairpin>" if next_char == ')' else "<internal_loop>")
             append_once(label)
             i = j
             continue
         if c == '(':
             append_once("<stem>")
             depth += 1
-        elif c == ')':
+        else:  # ')'
             append_once("<stem>")
             depth = max(depth - 1, 0)
         i += 1
@@ -109,23 +164,18 @@ def dotbracket_to_structural(dot_str: str) -> str:
     res.append("<end>")
     return "".join(res)
 
+# ---------- UI fn ----------
 def predict(seq: str):
     seq = (seq or "").strip().upper()
-    if not seq or not set(seq) <= {"A", "U", "C", "G"}:
+    if not seq or not set(seq) <= {"A","U","C","G"}:
         return "Please enter an RNA sequence (A/U/C/G)."
 
-    n = len(seq)
-    prompt = _make_prompt(seq)
-    text = _generate(prompt, max_new_tokens=n + 8)
+    db = _generate_db(seq)
 
-    db = _extract_dotbracket_strict(text, n)
+    # if model still produced all dots, try a naive hairpin so you get stems
+    if db.count("(") + db.count(")") == 0:
+        db = naive_hairpin(seq)
 
-    # fallback if model gave nothing meaningful
-    if not db:
-        return "<start><external_loop><end>"
-
-    # If output shorter or longer, pad/truncate safely
-    db = (db[:n] + "." * n)[:n]
     return dotbracket_to_structural(db)
 
 demo = gr.Interface(
@@ -138,4 +188,3 @@ demo = gr.Interface(
 
 if __name__ == "__main__":
     demo.launch()
-