changes

app.py

@@ -5,11 +5,6 @@ Upload a FASTA of homologous sequences. If more than 15 are given, 15 are
 randomly sampled and used to prompt AI4PD/ProtGPT3-MSA, which then generates
 N new family-consistent sequences in parallel. Generation stops at EOS or the
 <s> separator, so each output is exactly one new sequence.
-
-DEPLOY:
-  1. Create a new Space (SDK = Gradio) at huggingface.co/new-space.
-  2. Add this file as `app.py` and requirements.txt next to it.
-  3. It builds and goes live. Free CPU is enough; ZeroGPU/GPU is faster.
 """
 
 import random
@@ -26,7 +21,6 @@ MODEL_ID = "AI4PD/ProtGPT3-MSA"
 MAX_CONTEXT = 15           # the model is prompted with at most 15 homologs
 DIRECTION = "1"            # "1" = N-to-C (fixed; "2" would need reversed seqs)
 
-# bf16 on GPU; fall back to fp32 on the free CPU tier (bf16 CPU ops are spotty)
 DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
 tokenizer = AutoTokenizer.from_pretrained(
@@ -44,7 +38,6 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-# Stop on EOS *or* the sequence separator -> only one sequence is generated
 S_TOKEN_ID = tokenizer.convert_tokens_to_ids("<s>")
 EOS_IDS = [i for i in [tokenizer.eos_token_id, S_TOKEN_ID] if i is not None]
 PAD_ID = tokenizer.pad_token_id or tokenizer.eos_token_id or S_TOKEN_ID
@@ -75,9 +68,9 @@ def build_prompt(sequences: list, gap: bool = False, direction: str = "1") -> st
 
     tokens = ["<|bos|>", direction, gap_token]
     for seq in sequences:
-        tokens.append("<s>")                       # separator before each seq
+        tokens.append("<s>")
         tokens.extend(list(process_style(seq, gap=gap)))
-    tokens.append("<s>")                           # train-time separator before continuation
+    tokens.append("<s>")
     return " ".join(tokens)
 
 
@@ -120,7 +113,6 @@ def generate(fasta_file, n_sequences, max_new_tokens, temperature, top_p,
         random.seed(int(seed))
         torch.manual_seed(int(seed))
 
-    # Randomly sample up to 15 homologs for the prompt
     context = random.sample(seqs, MAX_CONTEXT) if len(seqs) > MAX_CONTEXT else seqs
 
     try:
@@ -137,12 +129,11 @@ def generate(fasta_file, n_sequences, max_new_tokens, temperature, top_p,
         do_sample=True,
         temperature=float(temperature),
         top_p=float(top_p),
-        eos_token_id=EOS_IDS,            # stop at EOS or <s> -> one sequence
+        eos_token_id=EOS_IDS,
         pad_token_id=PAD_ID,
         num_return_sequences=int(n_sequences),
     )
 
-    # Keep only the newly generated continuation (strip the prompt prefix)
     generated = output_ids[:, prompt_len:]
     records_out = []
     for i, row in enumerate(generated, 1):
@@ -164,34 +155,144 @@ def generate(fasta_file, n_sequences, max_new_tokens, temperature, top_p,
     return fasta_text, out_path
 
 
+# --------------------------------------------------------------------------- #
+# Header banner (inline SVG — ships inside the app, works in light/dark mode)
+# --------------------------------------------------------------------------- #
+def _banner_svg() -> str:
+    rng = random.Random(11)
+    bars = []
+    for i in range(6):
+        w = rng.randint(150, 225)
+        y = 100 + i * 20
+        op = 0.5 + i * 0.08
+        bars.append(
+            f'<rect x="60" y="{y}" width="{w}" height="11" rx="5.5" '
+            f'fill="url(#homo)" opacity="{op:.2f}"/>'
+        )
+    bars_svg = "\n".join(bars)
+    return f'''
+<div style="width:100%;max-width:980px;margin:0 auto 4px;">
+<svg viewBox="0 0 980 250" xmlns="http://www.w3.org/2000/svg"
+     style="width:100%;height:auto;display:block;border-radius:20px;">
+  <defs>
+    <linearGradient id="panel" x1="0" y1="0" x2="1" y2="1">
+      <stop offset="0" stop-color="#0f2742"/>
+      <stop offset="1" stop-color="#1c3a5e"/>
+    </linearGradient>
+    <linearGradient id="homo" x1="0" y1="0" x2="1" y2="0">
+      <stop offset="0" stop-color="#5b8fc7"/>
+      <stop offset="1" stop-color="#83b4dc"/>
+    </linearGradient>
+    <linearGradient id="gen" x1="0" y1="0" x2="1" y2="0">
+      <stop offset="0" stop-color="#34d29c"/>
+      <stop offset="1" stop-color="#7ff0cb"/>
+    </linearGradient>
+    <filter id="glow" x="-40%" y="-80%" width="180%" height="260%">
+      <feGaussianBlur stdDeviation="4" result="b"/>
+      <feMerge><feMergeNode in="b"/><feMergeNode in="SourceGraphic"/></feMerge>
+    </filter>
+    <marker id="arr" markerWidth="9" markerHeight="9" refX="5.5" refY="3" orient="auto">
+      <path d="M0,0 L6,3 L0,6 Z" fill="#86a5cb"/>
+    </marker>
+  </defs>
+
+  <rect x="0" y="0" width="980" height="250" rx="20" fill="url(#panel)"/>
+
+  <text x="60" y="46" font-family="system-ui,Segoe UI,Helvetica,Arial"
+        font-size="31" font-weight="700" fill="#eef4fb">ProtGPT3-MSA</text>
+  <text x="60" y="72" font-family="system-ui,Segoe UI,Helvetica,Arial"
+        font-size="14" fill="#9fbbd8">homolog-conditioned protein sequence generation</text>
+
+  {bars_svg}
+  <text x="60" y="236" font-family="system-ui,Segoe UI,Helvetica,Arial"
+        font-size="12" fill="#9fbbd8">up to 15 homologs</text>
+
+  <line x1="300" y1="151" x2="360" y2="151" stroke="#86a5cb"
+        stroke-width="2.5" marker-end="url(#arr)"/>
+
+  <rect x="384" y="119" width="64" height="64" rx="15"
+        fill="rgba(255,255,255,0.05)" stroke="#466a99" stroke-width="1.5"/>
+  <g stroke="#83b4dc" stroke-width="1.4" fill="#83b4dc">
+    <line x1="402" y1="135" x2="430" y2="151"/><line x1="402" y1="151" x2="430" y2="151"/>
+    <line x1="402" y1="167" x2="430" y2="151"/><line x1="402" y1="135" x2="430" y2="167"/>
+    <line x1="402" y1="167" x2="430" y2="135"/>
+    <circle cx="402" cy="135" r="3.4"/><circle cx="402" cy="151" r="3.4"/>
+    <circle cx="402" cy="167" r="3.4"/><circle cx="430" cy="151" r="3.8"/>
+  </g>
+  <text x="416" y="200" text-anchor="middle"
+        font-family="system-ui,Segoe UI,Helvetica,Arial"
+        font-size="12" fill="#9fbbd8">model</text>
+
+  <line x1="460" y1="151" x2="524" y2="151" stroke="#86a5cb"
+        stroke-width="2.5" marker-end="url(#arr)"/>
+
+  <rect x="552" y="130" width="222" height="13" rx="6.5" fill="url(#gen)" opacity="0.25"/>
+  <rect x="548" y="138" width="228" height="14" rx="7" fill="url(#gen)" opacity="0.45"/>
+  <rect x="544" y="147" width="234" height="17" rx="8.5" fill="url(#gen)" filter="url(#glow)"/>
+  <text x="544" y="186" font-family="system-ui,Segoe UI,Helvetica,Arial"
+        font-size="12" fill="#9fbbd8">N new sequences</text>
+</svg>
+</div>'''
+
+
+DESCRIPTION = """
+**ProtGPT3-MSA** is an autoregressive protein language model that can be prompted \
+with up to 15 homologous protein sequences to generate new, family-consistent \
+sequences. ([read the paper](https://www.biorxiv.org/content/10.64898/2026.06.04.730041v1))
+
+It operates in two modalities. In **unaligned** mode, homologs are supplied as raw \
+sequences and the model returns a plain sequence. In **aligned** mode, homologs are \
+supplied as a gapped multiple-sequence alignment and the model returns a new aligned \
+(gapped) sequence consistent with that alignment.
+
+**Settings**
+- **Upload a FASTA file** of homologous sequences. If it contains more than 15 \
+sequences, a random subset of 15 is sampled to build the prompt.
+- **Temperature** sets how stochastic generation is: raise it (e.g. > 1.0) for more \
+diverse, exploratory sequences, or lower it (e.g. < 0.7) to keep generation \
+conservative and close to the input family.
+- **Use aligned sequences** — tick this box to run in aligned mode, in which the model \
+generates aligned (gapped) sequences. Note that your uploaded homologs must then \
+already be aligned (equal length, with gap characters).
+
+**Recommendation.** ProtGPT3-MSA performs best when several homologs are provided. \
+With very few sequences (fewer than 5), generation quality may be limited.
+
+For full functionality and programmatic use, ProtGPT3-MSA is available on the \
+[Hugging Face Hub](https://huggingface.co/AI4PD/ProtGPT3-MSA).
+"""
+
+
 # --------------------------------------------------------------------------- #
 # UI
 # --------------------------------------------------------------------------- #
-demo = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.File(label="Homologs (FASTA)", file_types=[".fasta", ".fa", ".txt"],
-                type="filepath"),
-        gr.Slider(1, 100, value=20, step=1, label="N sequences to generate"),
-        gr.Slider(16, 1024, value=512, step=16, label="Max new tokens (per sequence)"),
-        gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature"),
-        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
-        gr.Checkbox(value=False, label="use aligned sequences"),
-        gr.Number(value=-1, label="Random seed (-1 = random)"),
-    ],
-    outputs=[
-        gr.Textbox(label="Generated sequences (FASTA)", lines=16),
-        gr.File(label="Download FASTA"),
-    ],
-    title="ProtGPT3-MSA — homolog-conditioned protein generator",
-    description=(
-        "Upload homologous protein sequences. Up to 15 are used to prompt "
-        "ProtGPT3-MSA, which generates N new family-consistent sequences. "
-        "More than 15 → a random subset is sampled. Tick **use aligned "
-        "sequences** to run in <gap> mode (requires equal-length, aligned input)."
-    ),
-    flagging_mode="never",
-)
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="teal"),
+               title="ProtGPT3-MSA") as demo:
+    gr.HTML(_banner_svg())
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            fasta = gr.File(label="Homologs (FASTA)",
+                            file_types=[".fasta", ".fa", ".txt"], type="filepath")
+            n_seq = gr.Slider(1, 100, value=20, step=1, label="N sequences to generate")
+            max_tok = gr.Slider(16, 1024, value=512, step=16,
+                                label="Max new tokens (per sequence)")
+            temp = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+            aligned = gr.Checkbox(value=False, label="use aligned sequences")
+            seed = gr.Number(value=-1, label="Random seed (-1 = random)")
+            btn = gr.Button("Generate", variant="primary")
+        with gr.Column(scale=1):
+            out_text = gr.Textbox(label="Generated sequences (FASTA)", lines=20)
+            out_file = gr.File(label="Download FASTA")
+
+    btn.click(
+        generate,
+        inputs=[fasta, n_seq, max_tok, temp, top_p, aligned, seed],
+        outputs=[out_text, out_file],
+        api_name="predict",
+    )
 
 if __name__ == "__main__":
     demo.launch()