Update app.py

app.py

@@ -1,10 +1,7 @@
 import os
-import re
-import math
 import torch
 import gradio as gr
 from functools import lru_cache
-from collections import deque
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -12,22 +9,8 @@ from transformers import (
     LogitsProcessorList,
 )
 
-# --------------------------- Config ---------------------------------
 MODEL_ID = os.getenv("MODEL_ID", "llm-rna-api-rmit/rna-structure-model")
 
-# sampling / decoding knobs (tune as you like)
-TEMPERATURE = float(os.getenv("TEMP", "0.8"))
-TOP_P = float(os.getenv("TOP_P", "0.9"))
-DOT_BIAS = float(os.getenv("DOT_BIAS", "0.8"))          # +logit added to '.'
-PAREN_RUN_PENALTY = float(os.getenv("PAREN_PEN", "0.5")) # -logit if long run
-PAREN_RUN_WINDOW = int(os.getenv("PAREN_WIN", "5"))      # lookback tokens
-
-# RNA constraints
-ALLOW_GU = os.getenv("ALLOW_GU", "1") != "0"
-MIN_LOOP_LEN = int(os.getenv("MIN_LOOP", "3"))  # minimum unpaired bases in loop
-
-# --------------------------------------------------------------------
-
 @lru_cache(maxsize=1)
 def _load_model_and_tokenizer():
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -40,25 +23,7 @@ def _load_model_and_tokenizer():
     model.eval()
     return tokenizer, model, device
 
-FEWSHOT = """Return ONLY the dot-bracket structure, one line, same length as RNA.
-RNA: GGGAAAUCCCU
-Dot-bracket: (((...)))).
-RNA: AUAUAUAU
-Dot-bracket: ........
-RNA: GGGAAACCC
-Dot-bracket: (((...)))
-RNA: GAAACUU
-Dot-bracket: (..())
-"""
-
-def _make_prompt(seq: str) -> str:
-    n = len(seq)
-    return (
-        FEWSHOT
-        + f"\nRNA: {seq}\n"
-        + f"Dot-bracket (exactly {n} characters using only '(' ')' '.'):\n"
-    )
-
+# --- Utility helpers ---
 def _char_token_id(tokenizer, ch: str) -> int:
     ids = tokenizer.encode(ch, add_special_tokens=False)
     for tid in ids:
@@ -67,261 +32,148 @@ def _char_token_id(tokenizer, ch: str) -> int:
     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}")
+    raise ValueError(f"Could not find token id for {ch}")
 
-# ---------- RNA pairing helpers ----------
-def _can_pair(a: str, b: str, allow_gu=True) -> bool:
-    if a == "A" and b == "U": return True
-    if a == "U" and b == "A": return True
-    if a == "G" and b == "C": return True
-    if a == "C" and b == "G": return True
-    if allow_gu and ((a == "G" and b == "U") or (a == "U" and b == "G")):
+def _can_pair(a, b, allow_gu=True):
+    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")]:
         return True
     return False
 
-def _precompute_can_open(seq: str, min_loop_len: int, allow_gu: bool):
-    """
-    can_open[i] = there exists j >= i + min_loop_len + 1 with pair(seq[i], seq[j])
-    """
-    n = len(seq)
-    can_open = [False] * n
-    # For speed, pre-index future positions by base
-    pos_by_base = {"A": [], "U": [], "G": [], "C": []}
-    for idx in range(n-1, -1, -1):
-        base = seq[idx]
-        # update future lists (right side of idx)
-        if idx + 1 < n:
-            pos_by_base[seq[idx+1]].append(idx+1)
-        # Check if any future partner exists with min loop
-        min_j = idx + min_loop_len + 1
-        ok = False
-        for b, lst in pos_by_base.items():
-            if any(j >= min_j and _can_pair(base, b, allow_gu) for j in lst):
-                ok = True
+def _precompute_can_open(seq, min_loop=3, allow_gu=True):
+    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
                 break
-        can_open[idx] = ok
-    return can_open
-
-# --------- Top-p sampling ----------
-def _top_p_sample_from_logits(logits: torch.Tensor, top_p: float, temperature: float) -> torch.Tensor:
-    if temperature <= 0:
-        # fall back to greedy
-        return torch.argmax(logits, dim=-1)
-    logits = logits / temperature
-    probs = torch.softmax(logits, dim=-1)
-
-    # sort
-    sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
-    cumulative = torch.cumsum(sorted_probs, dim=-1)
-
-    # mask tokens beyond top_p
-    cutoff = (cumulative > top_p).float()
-    # ensure at least one token kept
-    cutoff[..., 0] = 0.0
-    keep = 1.0 - cutoff
-
-    filtered_probs = sorted_probs * keep
-    # renormalize
-    filtered_probs = filtered_probs / filtered_probs.sum(dim=-1, keepdim=True)
+    return can
 
-    # sample
-    next_sorted = torch.multinomial(filtered_probs, num_samples=1)
-    next_ids = sorted_idx.gather(-1, next_sorted)
-    return next_ids.squeeze(-1)
-
-# --- Finite-State logits processor with RNA constraints ---
+# --- constrained processor ---
 class BalancedParenProcessor(LogitsProcessor):
-    """
-    Keeps output a valid dot-bracket string length N and adds biases:
-      - Balance constraints (no ')' at depth 0; enough room to close)
-      - Force close-all when remaining == depth
-      - Forbid '(' if no feasible partner ahead (RNA pairing + min loop)
-      - Dot bias (+bias to dot logit)
-      - Penalize long runs of parentheses
-    """
-    def __init__(self, lp_id: int, rp_id: int, dot_id: int, total_len: int,
-                 can_open, paren_run_window=5, paren_run_penalty=0.5, dot_bias=0.8):
-        self.lp_id = int(lp_id)
-        self.rp_id = int(rp_id)
-        self.dot_id = int(dot_id)
-        self.total_len = int(total_len)
-
+    def __init__(self, lp_id, rp_id, dot_id, total_len, can_open,
+                 dot_bias=0.8, paren_penalty=0.5, 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.recent = deque(maxlen=paren_run_window)
-        self.paren_run_penalty = float(paren_run_penalty)
-        self.dot_bias = float(dot_bias)
-        self.can_open = can_open  # list[bool] length N
-
-    def update_with_chosen(self, token_id: int):
-        # called after each step by the generator
-        self.recent.append(token_id)
-        if token_id == self.lp_id:
-            self.depth += 1
-        elif token_id == self.rp_id:
-            self.depth = max(0, self.depth - 1)
-        self.step += 1
-
-    def _recent_is_paren_run(self):
-        if not self.recent:
-            return False
-        return all(t in (self.lp_id, self.rp_id) for t in self.recent)
+        self.history=[]
+        self.can_open = can_open
+        self.dot_bias=dot_bias
+        self.paren_penalty=paren_penalty
+        self.window=window
 
     def __call__(self, input_ids, scores):
-        # scores: (1, vocab)
-        mask = torch.full_like(scores, float("-inf"))
-
-        remaining = self.total_len - self.step
-        allowed = []
-
-        must_close_all = (remaining == self.depth and self.depth > 0)
-        pos = self.step
-
-        if must_close_all:
-            allowed = [self.rp_id]
+        mask=torch.full_like(scores,float("-inf"))
+        remaining=self.total_len-self.step
+        allowed=[]
+        must_close=(remaining==self.depth and self.depth>0)
+        pos=self.step
+        if must_close:
+            allowed=[self.rp_id]
         else:
-            # Allow ')' only if inside a stem
-            if self.depth > 0:
+            if self.depth>0:
                 allowed.append(self.rp_id)
-
-            # Allow '(' only if (a) room to close by end and (b) feasible partner ahead
-            if remaining - 1 > self.depth and pos < self.total_len and self.can_open[pos]:
+            if remaining-1>self.depth and pos<len(self.can_open) and self.can_open[pos]:
                 allowed.append(self.lp_id)
-
-            # '.' generally allowed
             allowed.append(self.dot_id)
-
-        mask[:, allowed] = 0.0
-        scores = scores + mask
-
-        # add dot bias
-        scores[..., self.dot_id] = scores[..., self.dot_id] + self.dot_bias
-
-        # penalize long paren run
-        if self._recent_is_paren_run():
-            scores[..., self.lp_id] = scores[..., self.lp_id] - self.paren_run_penalty
-            scores[..., self.rp_id] = scores[..., self.rp_id] - self.paren_run_penalty
-
+        mask[:,allowed]=0.0
+        scores=scores+mask
+        scores[:,self.dot_id]+=self.dot_bias
+        if len(self.history)>=self.window and all(t in (self.lp_id,self.rp_id) for t in self.history[-self.window:]):
+            scores[:,self.lp_id]-=self.paren_penalty
+            scores[:,self.rp_id]-=self.paren_penalty
         return scores
 
-# --- Generate exactly n chars using constrained decoding + sampling ---
-def _generate_db(seq: str) -> str:
-    tokenizer, model, device = _load_model_and_tokenizer()
-    n = len(seq)
-    prompt = _make_prompt(seq)
-
-    lp_id = _char_token_id(tokenizer, "(")
-    rp_id = _char_token_id(tokenizer, ")")
-    dot_id = _char_token_id(tokenizer, ".")
-
-    can_open = _precompute_can_open(seq, MIN_LOOP_LEN, ALLOW_GU)
-
-    processor = BalancedParenProcessor(
-        lp_id, rp_id, dot_id, n,
-        can_open=can_open,
-        paren_run_window=PAREN_RUN_WINDOW,
-        paren_run_penalty=PAREN_RUN_PENALTY,
-        dot_bias=DOT_BIAS,
-    )
-    processors = LogitsProcessorList([processor])
-
-    with torch.inference_mode():
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
-
-        generated = []
-        cur_input = inputs["input_ids"]
-        for _ in range(n):
-            outputs = model(cur_input)
-            logits = outputs.logits[:, -1, :]  # (1, vocab)
-            for p in processors:
-                logits = p(cur_input, logits)
-
-            next_id = _top_p_sample_from_logits(logits, TOP_P, TEMPERATURE).unsqueeze(0)
-            token_id = next_id.item()
-            generated.append(token_id)
-
-            processor.update_with_chosen(token_id)
-
-            cur_input = torch.cat([cur_input, next_id.unsqueeze(0)], dim=1)
-
-    text = tokenizer.decode(generated, skip_special_tokens=True)
-    db = "".join(c for c in text if c in "().")[:n]
-    if len(db) != n:
-        db = (db + "." * n)[:n]
+    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)
+    cum=torch.cumsum(sorted_probs,dim=-1)
+    mask=cum>top_p
+    mask[...,0]=False
+    sorted_probs[mask]=0
+    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)
+    proc=BalancedParenProcessor(lp,rp,dot,n,can)
+    procs=LogitsProcessorList([proc])
+    inputs=tok(prompt,return_tensors="pt").to(device)
+    cur=inputs["input_ids"]
+    generated=[]
+    for _ in range(n):
+        out=model(cur)
+        logits=out.logits[:,-1,:]
+        for p in procs:
+            logits=p(cur,logits)
+        next_id=_top_p_sample(logits,0.9,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)
+    text=tok.decode(generated,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
 
-# --- Dot-bracket -> structural elements ---
-def dotbracket_to_structural(dot_str: str) -> str:
-    if not dot_str:
-        return "<start><external_loop><end>"
-    res = ["<start>"]
-    depth = 0
-    i, n = 0, len(dot_str)
-
-    def append_once(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
-            next_char = dot_str[j] if j < n else None
-            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
-        else:  # ')'
-            append_once("<stem>")
-            depth = max(depth - 1, 0)
-        i += 1
-
+# --- 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
     res.append("<end>")
     return "".join(res)
 
-# --- Gradio handler ---
-def predict(seq: str):
-    seq = (seq or "").strip().upper()
-    if not seq or not set(seq) <= {"A","U","C","G"}:
+# --- 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)
+    db=_generate_db(seq)
     return dotbracket_to_structural(db)
 
-# UI with a few knobs exposed
-with gr.Blocks(title="RNA Structure Predictor") as demo:
-    gr.Markdown("### RNA Structure Predictor\nOutputs structural-element notation: `<start>`, `<stem>`, `<hairpin>`, `<internal_loop>`, `<external_loop>`, `<end>`.")
-    with gr.Row():
-        seq_in = gr.Textbox(lines=4, label="RNA Sequence (A/U/C/G)", value="GGGAAUCC")
-        out = gr.Textbox(lines=6, label="Predicted Structural Elements")
-    with gr.Row():
-        t = gr.Slider(0.1, 1.5, value=TEMPERATURE, step=0.05, label="Temperature")
-        p = gr.Slider(0.5, 1.0, value=TOP_P, step=0.01, label="Top-p")
-        dbias = gr.Slider(0.0, 2.0, value=DOT_BIAS, step=0.05, label="Dot bias (+logit)")
-        looplen = gr.Slider(0, 5, value=MIN_LOOP_LEN, step=1, label="Min loop length")
-        wobble = gr.Checkbox(value=ALLOW_GU, label="Allow GU wobble")
-    btn = gr.Button("Submit", variant="primary")
-
-    def _predict_with_knobs(seq, temperature, topp, dot_bias, min_loop, allow_gu):
-        global TEMPERATURE, TOP_P, DOT_BIAS, MIN_LOOP_LEN, ALLOW_GU
-        TEMPERATURE = float(temperature)
-        TOP_P = float(topp)
-        DOT_BIAS = float(dot_bias)
-        MIN_LOOP_LEN = int(min_loop)
-        # this affects precompute_can_open on next call
-        ALLOW_GU = bool(allow_gu)
-        return predict(seq)
-
-    btn.click(
-        _predict_with_knobs,
-        inputs=[seq_in, t, p, dbias, looplen, wobble],
-        outputs=[out],
-    )
+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"),
+    title="RNA Structure Predictor",
+    description="Outputs <start>, <stem>, <hairpin>, <internal_loop>, <external_loop>, <end>."
+)
 
-if __name__ == "__main__":
+if __name__=="__main__":
     demo.launch()