app.py

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("""
        <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;">
                    &nbsp;Bidirectional Protein-Text Generation
                </span>
            </h1>
            <p style="margin:0.4rem 0 0 0; color:#888; font-size:0.9rem;">
                ⚡ Running on CPU &mdash; each generation takes ~30&ndash;60 seconds
            </p>
        </div>
        """)

        # ==================== TAB 1: Protein -> Text ====================
        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,
            )

        # ==================== TAB 2: Text -> Protein ====================
        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.'
            )

            # ---- Structured Prompt Builder ----
            gr.Markdown("### \U0001f3d7 Structured Prompt Builder")

            # Common fields (7)
            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.")

            # 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('<div class="footer">Ouroboros &mdash; Dual T5 Model</div>')

    return demo


if __name__ == "__main__":
    demo = create_demo()
    demo.queue(max_size=5).launch()