| import gradio as gr |
| import torch |
| import os |
| import json |
| import random |
|
|
| from ouroboros_model import OuroborosModel |
|
|
| MODEL = None |
| SAMPLES = None |
|
|
| MODEL_REPO = os.environ.get("OUROBOROS_MODEL_REPO", "REXWind/ouroboros-weights") |
| CHECKPOINT_PATH = os.environ.get("CHECKPOINT_PATH", |
| os.path.join(os.path.dirname(__file__), "dual_model.pt")) |
| SAMPLE_FILE = os.path.join(os.path.dirname(__file__), "100sample.json") |
|
|
| ORDER_LIST = [ |
| 'similarity', 'function', 'subunit', 'catalytic activity', 'pathway', |
| 'cofactor', 'tissue specificity', 'domain', 'subcellular location', 'PTM', |
| 'miscellaneous', 'induction', 'disruption phenotype', 'activity regulation', |
| 'developmental stage', 'sequence caution', 'caution', |
| 'biophysicochemical properties', 'disease', 'alternative products', |
| 'online information', 'biotechnology', 'polymorphism', 'mass spectrometry', |
| 'allergen', 'toxic dose', 'RNA editing', 'pharmaceutical', 'interaction' |
| ] |
|
|
|
|
| def load_model(): |
| global MODEL |
| if MODEL is not None: |
| return |
| print("Loading Ouroboros model...") |
| device = "cpu" |
| if not os.path.exists(CHECKPOINT_PATH) and MODEL_REPO: |
| from huggingface_hub import hf_hub_download |
| print(f"Downloading checkpoint from {MODEL_REPO}...") |
| checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename="dual_model.pt") |
| else: |
| checkpoint_path = CHECKPOINT_PATH |
| if not os.path.exists(checkpoint_path): |
| raise FileNotFoundError(f"Checkpoint not found at {checkpoint_path}") |
| MODEL = OuroborosModel( |
| t5config_path=os.path.join(os.path.dirname(__file__), "t5config_large.json"), |
| text_share_param=True, protein_share_param=True, share_fc=True, |
| use_mean_pooling=False, |
| ) |
| MODEL.load_checkpoint(checkpoint_path) |
| MODEL.to(device) |
| MODEL.eval() |
| print(f"Model loaded on {device}") |
|
|
|
|
| def load_samples(): |
| global SAMPLES |
| if SAMPLES is not None: |
| return |
| if os.path.exists(SAMPLE_FILE): |
| with open(SAMPLE_FILE, 'r') as f: |
| SAMPLES = json.load(f) |
| print(f"Loaded {len(SAMPLES)} samples") |
| else: |
| SAMPLES = [] |
| print("No sample file found") |
|
|
|
|
| def protein_to_text(protein_seq, temperature, top_k, top_p, num_beams, use_greedy): |
| if MODEL is None: |
| return "Model not loaded yet, please wait..." |
| if not protein_seq or not protein_seq.strip(): |
| return "Please enter a protein sequence." |
| seq = protein_seq.strip().replace(" ", "") |
| seq = " ".join(seq) |
| try: |
| texts = MODEL.protein_to_text( |
| [seq], use_greedy=use_greedy, temperature=temperature, |
| top_k=int(top_k), top_p=top_p, num_beams=int(num_beams), |
| ) |
| return texts[0] |
| except Exception as e: |
| return f"Error: {e}" |
|
|
|
|
| def text_to_protein(text_desc, num_return, temperature, top_k, top_p, num_beams, use_greedy): |
| if MODEL is None: |
| return "Model not loaded yet, please wait..." |
| if not text_desc or not text_desc.strip(): |
| return "Please enter a text description." |
| try: |
| sequences = MODEL.text_to_protein( |
| [text_desc.strip()], use_greedy=use_greedy, temperature=temperature, |
| top_k=int(top_k), top_p=top_p, num_beams=int(num_beams), |
| num_return_sequences=int(num_return), |
| ) |
| if len(sequences) == 1: |
| return sequences[0] |
| return "\n\n".join(f"[{i + 1}] {s}" for i, s in enumerate(sequences)) |
| except Exception as e: |
| return f"Error: {e}" |
|
|
|
|
| def assemble_prompt(*field_values): |
| """Concatenate all field values in ORDER_LIST sequence.""" |
| parts = [] |
| for i, key in enumerate(ORDER_LIST): |
| if i < len(field_values): |
| val = field_values[i].strip() |
| if val: |
| parts.append(val) |
| return "".join(parts) |
|
|
|
|
| def random_example(): |
| """Pick a random sample; return all field values + assembled prompt.""" |
| load_samples() |
| if not SAMPLES: |
| empty_fields = [""] * len(ORDER_LIST) |
| return empty_fields + ["No samples available"] |
|
|
| sample = random.choice(SAMPLES) |
| values = [] |
| for key in ORDER_LIST: |
| raw = sample.get(key, "") |
| if isinstance(raw, list): |
| values.append(raw[0] if raw else "") |
| else: |
| values.append(str(raw)) |
|
|
| parts = [] |
| for v in values: |
| if v.strip(): |
| parts.append(v.strip()) |
| prompt = "".join(parts) |
| return values + [prompt] |
|
|
|
|
| def create_demo(): |
| load_model() |
|
|
| css = """ |
| /* ===== Color Palette ===== |
| Protein: accent=#C0392B border=#E6B0AA bg=#FDEDEC |
| Text: accent=#2471A3 border=#A9CCE3 bg=#EBF5FB |
| ====================================== */ |
| |
| /* ---- Header ---- */ |
| .header-box { |
| background: linear-gradient(135deg, #EBF5FB 0%, #FDEDEC 100%); |
| border-radius: 16px; |
| padding: 1.8rem 2rem; |
| margin-bottom: 1.2rem; |
| border: 1px solid #D5D8DC; |
| } |
| .header-box h1 { margin-bottom: 0.3rem; } |
| |
| /* ---- Protein (red/pink) inputs ---- */ |
| .protein-input textarea { |
| border: 2px solid #E6B0AA !important; |
| border-radius: 8px !important; |
| font-family: 'Courier New', monospace !important; |
| font-size: 15px !important; |
| } |
| .protein-input:focus-within textarea { |
| border-color: #E74C3C !important; |
| box-shadow: 0 0 0 2px rgba(231, 76, 60, 0.12) !important; |
| } |
| |
| /* ---- Text (blue) inputs ---- */ |
| .text-input textarea { |
| border: 2px solid #A9CCE3 !important; |
| border-radius: 8px !important; |
| } |
| .text-input:focus-within textarea { |
| border-color: #2471A3 !important; |
| box-shadow: 0 0 0 2px rgba(36, 113, 163, 0.12) !important; |
| } |
| |
| /* ---- Output boxes ---- */ |
| .protein-output textarea { |
| border-left: 4px solid #E74C3C !important; |
| border-radius: 0 8px 8px 0 !important; |
| font-family: 'Courier New', monospace !important; |
| font-size: 14px !important; |
| } |
| .text-output textarea { |
| border-left: 4px solid #2471A3 !important; |
| border-radius: 0 8px 8px 0 !important; |
| } |
| |
| /* ---- Buttons ---- */ |
| .btn-protein { |
| background: #E74C3C !important; |
| color: white !important; |
| border: none !important; |
| font-weight: 600 !important; |
| } |
| .btn-protein:hover { background: #CB4335 !important; } |
| |
| .btn-text { |
| background: #2471A3 !important; |
| color: white !important; |
| border: none !important; |
| font-weight: 600 !important; |
| } |
| .btn-text:hover { background: #1A5276 !important; } |
| |
| .btn-example { |
| background: #FFFFFF !important; |
| color: #2471A3 !important; |
| border: 2px solid #A9CCE3 !important; |
| font-weight: 500 !important; |
| } |
| .btn-example:hover { background: #EBF5FB !important; } |
| |
| /* ---- Assembled prompt ---- */ |
| .assembled-prompt textarea { |
| border: 2px dashed #A9CCE3 !important; |
| background: #FAFAFA !important; |
| font-size: 13px !important; |
| border-radius: 8px !important; |
| } |
| |
| /* ---- Group sections ---- */ |
| .section-group { |
| border: 1px solid #E5E7E9; |
| border-radius: 10px; |
| padding: 1rem; |
| margin-bottom: 0.8rem; |
| background: #FBFCFC; |
| } |
| |
| /* ---- Field textboxes ---- */ |
| .field-textbox textarea { |
| font-size: 12px !important; |
| border-radius: 6px !important; |
| border: 1px solid #D5D8DC !important; |
| } |
| .field-textbox:focus-within textarea { |
| border-color: #2471A3 !important; |
| } |
| |
| /* ---- Slider colors ---- */ |
| input[type="range"] { accent-color: #2471A3; } |
| |
| /* ---- Accordion ---- */ |
| .acc-more-fields { border: 1px solid #D5D8DC; border-radius: 8px; } |
| |
| /* ---- Footer ---- */ |
| .footer { text-align: center; color: #999; font-size: 13px; margin-top: 2rem; } |
| """ |
|
|
| with gr.Blocks(title="Ouroboros", css=css, theme=gr.themes.Soft()) as demo: |
| |
| gr.HTML(""" |
| <div class="header-box"> |
| <h1 style="margin:0; font-size:1.8rem;"> |
| <span style="color:#E74C3C;">Ouroboros</span> |
| <span style="color:#666; font-size:1.2rem; font-weight:400;"> |
| Bidirectional Protein-Text Generation |
| </span> |
| </h1> |
| <p style="margin:0.4rem 0 0 0; color:#888; font-size:0.9rem;"> |
| ⚡ Running on CPU — each generation takes ~30–60 seconds |
| </p> |
| </div> |
| """) |
|
|
| |
| with gr.Tab("\U0001f9ec Protein \u2192 Text"): |
| gr.Markdown( |
| '<span style="color:#E74C3C; font-weight:600;">Input</span> ' |
| 'a protein sequence to generate its functional description.' |
| ) |
| protein_input = gr.Textbox( |
| label="Protein Sequence", |
| placeholder="MAKNVLAVTGSSDGFGAVAAALLGADVAIVGTGRPKALADLARELGAR...", |
| lines=4, elem_classes="protein-input", |
| ) |
| p2t_output = gr.Textbox( |
| label="Generated Description", |
| lines=5, interactive=False, elem_classes="text-output", |
| ) |
| p2t_btn = gr.Button( |
| "Generate Description", variant="primary", elem_classes="btn-protein", |
| ) |
|
|
| with gr.Accordion("\u2699 Advanced Settings", open=False): |
| with gr.Row(): |
| p2t_temp = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature") |
| p2t_topk = gr.Slider(0, 100, value=40, step=1, label="Top-K") |
| p2t_topp = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-P") |
| p2t_beams = gr.Slider(1, 10, value=4, step=1, label="Num Beams") |
| p2t_greedy = gr.Checkbox(value=False, label="Greedy Decoding") |
|
|
| p2t_btn.click( |
| protein_to_text, |
| inputs=[protein_input, p2t_temp, p2t_topk, p2t_topp, p2t_beams, p2t_greedy], |
| outputs=p2t_output, |
| ) |
|
|
| |
| with gr.Tab("\U0001f4dd Text \u2192 Protein"): |
| gr.Markdown( |
| '<span style="color:#2471A3; font-weight:600;">Describe</span> ' |
| 'a protein\'s function to generate candidate sequences. ' |
| 'Fill in the structured fields below, or use <b>Random Example</b> for inspiration.' |
| ) |
|
|
| |
| gr.Markdown("### \U0001f3d7 Structured Prompt Builder") |
|
|
| |
| with gr.Group(elem_classes="section-group"): |
| gr.Markdown('<p style="color:#2471A3; font-weight:600; margin:0 0 0.5rem 0;">Common Annotations</p>') |
| with gr.Row(): |
| f_similarity = gr.Textbox(label="Similarity", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. Belongs to the MurCDEF family.") |
| f_function = gr.Textbox(label="Function", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. Cell wall formation.") |
| f_catalytic = gr.Textbox(label="Catalytic Activity", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. ATP + H2O = ADP + phosphate") |
| f_subunit = gr.Textbox(label="Subunit", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. Homodimer.") |
| with gr.Row(): |
| f_cofactor = gr.Textbox(label="Cofactor", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. Binds 1 zinc ion.") |
| f_pathway = gr.Textbox(label="Pathway", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. Peptidoglycan biosynthesis.") |
| f_subcell_loc = gr.Textbox(label="Subcellular Location", lines=1, elem_classes="field-textbox", |
| placeholder="e.g. Cytoplasm.") |
|
|
| |
| with gr.Accordion("\U0001f4c2 More Annotation Fields", open=False, elem_classes="acc-more-fields"): |
| with gr.Row(): |
| f_tissue = gr.Textbox(label="Tissue Specificity", lines=1, elem_classes="field-textbox") |
| f_domain = gr.Textbox(label="Domain", lines=1, elem_classes="field-textbox") |
| f_ptm = gr.Textbox(label="PTM", lines=1, elem_classes="field-textbox") |
| f_misc = gr.Textbox(label="Miscellaneous", lines=1, elem_classes="field-textbox") |
| with gr.Row(): |
| f_induction = gr.Textbox(label="Induction", lines=1, elem_classes="field-textbox") |
| f_disruption = gr.Textbox(label="Disruption Phenotype", lines=1, elem_classes="field-textbox") |
| f_activity_reg = gr.Textbox(label="Activity Regulation", lines=1, elem_classes="field-textbox") |
| f_dev_stage = gr.Textbox(label="Developmental Stage", lines=1, elem_classes="field-textbox") |
| with gr.Row(): |
| f_seq_caution = gr.Textbox(label="Sequence Caution", lines=1, elem_classes="field-textbox") |
| f_caution = gr.Textbox(label="Caution", lines=1, elem_classes="field-textbox") |
| f_biophys = gr.Textbox(label="Biophysicochemical Properties", lines=1, elem_classes="field-textbox") |
| f_disease = gr.Textbox(label="Disease", lines=1, elem_classes="field-textbox") |
| with gr.Row(): |
| f_alt_products = gr.Textbox(label="Alternative Products", lines=1, elem_classes="field-textbox") |
| f_online_info = gr.Textbox(label="Online Information", lines=1, elem_classes="field-textbox") |
| f_biotech = gr.Textbox(label="Biotechnology", lines=1, elem_classes="field-textbox") |
| f_polymorphism = gr.Textbox(label="Polymorphism", lines=1, elem_classes="field-textbox") |
| with gr.Row(): |
| f_mass_spec = gr.Textbox(label="Mass Spectrometry", lines=1, elem_classes="field-textbox") |
| f_allergen = gr.Textbox(label="Allergen", lines=1, elem_classes="field-textbox") |
| f_toxic_dose = gr.Textbox(label="Toxic Dose", lines=1, elem_classes="field-textbox") |
| f_rna_editing = gr.Textbox(label="RNA Editing", lines=1, elem_classes="field-textbox") |
| with gr.Row(): |
| f_pharma = gr.Textbox(label="Pharmaceutical", lines=1, elem_classes="field-textbox") |
| f_interaction = gr.Textbox(label="Interaction", lines=1, elem_classes="field-textbox") |
|
|
| |
| with gr.Row(): |
| example_btn = gr.Button( |
| "\U0001f3b2 Random Example", variant="secondary", elem_classes="btn-example", |
| ) |
| build_btn = gr.Button( |
| "\u2699 Build Prompt", variant="primary", elem_classes="btn-text", |
| ) |
|
|
| |
| assembled_prompt = gr.Textbox( |
| label="\U0001f4cb Assembled Prompt (editable)", |
| placeholder="Click 'Build Prompt' or 'Random Example', or type/paste directly...", |
| lines=4, elem_classes="assembled-prompt", |
| ) |
|
|
| |
| with gr.Accordion("\u270f Direct Input (skip builder)", open=False): |
| direct_input = gr.Textbox( |
| label="Paste a full text prompt directly", |
| placeholder="e.g. Belongs to the MurCDEF family.Cell wall formation.UDP-N-acetyl-alpha-D-muramate + L-alanine + ATP = ...", |
| lines=3, |
| ) |
| direct_btn = gr.Button("Use This as Prompt", variant="secondary") |
| direct_btn.click(lambda x: x, inputs=[direct_input], outputs=[assembled_prompt]) |
|
|
| |
| t2p_output = gr.Textbox( |
| label="\U0001f9ec Generated Protein Sequence(s)", |
| lines=6, interactive=False, elem_classes="protein-output", |
| ) |
| t2p_btn = gr.Button( |
| "\u26a1 Generate Sequence", variant="primary", elem_classes="btn-protein", |
| ) |
|
|
| |
| with gr.Accordion("\u2699 Advanced Generation Settings", open=False): |
| with gr.Row(): |
| num_return = gr.Slider(1, 5, value=1, step=1, label="Num Return Sequences") |
| with gr.Row(): |
| t2p_temp = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature") |
| t2p_topk = gr.Slider(0, 100, value=40, step=1, label="Top-K") |
| t2p_topp = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-P") |
| t2p_beams = gr.Slider(1, 10, value=4, step=1, label="Num Beams") |
| t2p_greedy = gr.Checkbox(value=False, label="Greedy Decoding") |
|
|
| |
| all_fields = [ |
| f_similarity, f_function, f_subunit, f_catalytic, f_pathway, |
| f_cofactor, f_tissue, f_domain, f_subcell_loc, f_ptm, |
| f_misc, f_induction, f_disruption, f_activity_reg, |
| f_dev_stage, f_seq_caution, f_caution, |
| f_biophys, f_disease, f_alt_products, |
| f_online_info, f_biotech, f_polymorphism, f_mass_spec, |
| f_allergen, f_toxic_dose, f_rna_editing, f_pharma, |
| f_interaction, |
| ] |
| assert len(all_fields) == len(ORDER_LIST), \ |
| f"Field count mismatch: {len(all_fields)} vs {len(ORDER_LIST)}" |
|
|
| |
| build_btn.click(fn=assemble_prompt, inputs=all_fields, outputs=[assembled_prompt]) |
| example_btn.click(fn=random_example, inputs=[], outputs=all_fields + [assembled_prompt]) |
| t2p_btn.click( |
| text_to_protein, |
| inputs=[assembled_prompt, num_return, t2p_temp, t2p_topk, t2p_topp, t2p_beams, t2p_greedy], |
| outputs=t2p_output, |
| ) |
|
|
| |
| gr.HTML('<div class="footer">Ouroboros — Dual T5 Model</div>') |
|
|
| return demo |
|
|
|
|
| if __name__ == "__main__": |
| demo = create_demo() |
| demo.queue(max_size=5).launch() |
|
|
