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: # ---- Header ---- gr.HTML("""

Ouroboros  Bidirectional Protein-Text Generation

⚡ Running on CPU — each generation takes ~30–60 seconds

""") # ==================== TAB 1: Protein -> Text ==================== with gr.Tab("\U0001f9ec Protein \u2192 Text"): gr.Markdown( 'Input ' '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, ) # ==================== TAB 2: Text -> Protein ==================== with gr.Tab("\U0001f4dd Text \u2192 Protein"): gr.Markdown( 'Describe ' 'a protein\'s function to generate candidate sequences. ' 'Fill in the structured fields below, or use Random Example for inspiration.' ) # ---- Structured Prompt Builder ---- gr.Markdown("### \U0001f3d7 Structured Prompt Builder") # Common fields (7) with gr.Group(elem_classes="section-group"): gr.Markdown('

Common Annotations

') 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.") # More fields (22) 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") # Buttons 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 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", ) # Direct Input 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]) # Generated protein output 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", ) # Advanced settings 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 in ORDER_LIST order (29 items) 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)}" # Event Bindings 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, ) # ---- Footer ---- gr.HTML('
Ouroboros — Dual T5 Model
') return demo if __name__ == "__main__": demo = create_demo() demo.queue(max_size=5).launch()