File size: 4,509 Bytes
1e62ceb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import gradio as gr
import os
import pandas as pd
import tempfile
from pathlib import Path
from inference_wrapper import RNAFoldingPredictor
from huggingface_hub import hf_hub_download

# 🔽 Download the model weights from Hugging Face Hub
ckpt_file = hf_hub_download(
    repo_id="ptope/rna-model-weights",
    filename="best_val_model.pt"
)

# ✅ Set up model and static directory
MODEL_PATH = ckpt_file
STATIC_DIR = Path("static")
os.makedirs(STATIC_DIR, exist_ok=True)
predictor = RNAFoldingPredictor(MODEL_PATH)

# ✅ Load test sequences and map sequences to target IDs
test_sequences = pd.read_csv("test_sequences.csv")
test_seq_list = test_sequences["sequence"].dropna().unique().tolist()
sequence_to_target_id = {
    row["sequence"].strip().upper(): row["target_id"]
    for _, row in test_sequences.iterrows()
}

# ✅ Format predicted structure into standard PDB format
def create_pdb_from_prediction(prediction_df):
    pdb_lines = ["HEADER    RNA STRUCTURE PREDICTION"]
    atom_index = 1
    prev_c1_index = None
    for _, row in prediction_df.iterrows():
        resname = row['resname']
        resid = row['resid']
        x, y, z = row['x_1'], row['y_1'], row['z_1']
        c1_index = atom_index
        pdb_lines.append(f"ATOM  {atom_index:5d}  C1'  {resname} A{resid:4d}    {x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           C")
        atom_index += 1
        base_offset = {'A': (1.5, 0.0, 0.5), 'C': (1.2, 0.3, 0.3), 'G': (1.7, -0.2, 0.7), 'U': (1.0, 0.5, 0.0)}.get(resname, (1.5, 0.0, 0.0))
        base_x, base_y, base_z = x + base_offset[0], y + base_offset[1], z + base_offset[2]
        base_index = atom_index
        pdb_lines.append(f"ATOM  {atom_index:5d}  N9   {resname} A{resid:4d}    {base_x:8.3f}{base_y:8.3f}{base_z:8.3f}  1.00  0.00           N")
        atom_index += 1
        pdb_lines.append(f"CONECT{c1_index:5d}{base_index:5d}")
        if prev_c1_index is not None:
            pdb_lines.append(f"CONECT{prev_c1_index:5d}{c1_index:5d}")
        prev_c1_index = c1_index
    pdb_lines.append("END")
    return "\n".join(pdb_lines)

# ✅ This function handles GPT model prediction and downloadable .pdb
def generate_and_return_file(sequence, description=""):
    sequence = sequence.strip().upper()
    if not sequence or not all(b in "ACGU" for b in sequence):
        return "Invalid input", None

    df = predictor.predict(sequence, description)
    pdb_text = create_pdb_from_prediction(df)
    filename = f"{next(tempfile._get_candidate_names())}.pdb"
    path = STATIC_DIR / filename

    with open(path, "w") as f:
        f.write(pdb_text)

    return "Download ready", str(path)

# ✅ Load precomputed .pdb from /pdbs folder for Mol* viewer
def serve_precomputed_pdb(sequence, description=""):
    sequence = sequence.strip().upper()
    if sequence not in sequence_to_target_id:
        return ""

    target_id = sequence_to_target_id[sequence]
    filename = f"{target_id}.pdb"

    # ✅ Here's the FIX: point viewer to /pdbs folder

    pdb_url = f"https://ptope--rna-folding-demo.hf.space/pdbs/{filename}"
    iframe_url = f"https://molstar.org/viewer/?loadFromUrl={pdb_url}"

    # Check
    print(f"[DEBUG] Final URL to check: {pdb_url}")
    print(f"[DEBUG] Final URL to check: {iframe_url}")
    
    iframe_html = f'<iframe src="{iframe_url}" width="100%" height="600px" style="border:1px solid #ccc;"></iframe>'
    return iframe_html

# ✅ Gradio app layout
with gr.Blocks() as demo:
    gr.Markdown("## 🧬 RNA 3D Viewer: Precomputed Visualization + Predictive Download")

    with gr.Row():
        with gr.Column():
            dropdown = gr.Dropdown(choices=test_seq_list, label="Select a Test Sequence")
            seq_input = gr.Textbox(label="RNA Sequence", lines=4)
            desc_input = gr.Textbox(label="Description (optional)", lines=1)
            dropdown.change(fn=lambda s: s, inputs=dropdown, outputs=seq_input)
            submit_btn = gr.Button("Submit")

        with gr.Column():
            status_output = gr.Textbox(label="Status")
            file_output = gr.File(label="Download .pdb")

    viewer_html = gr.HTML(label="Mol* Viewer")

    # ✅ Submit triggers BOTH file generation and structure visualization
    submit_btn.click(
        fn=generate_and_return_file,
        inputs=[seq_input, desc_input],
        outputs=[status_output, file_output]
    )

    submit_btn.click(
        fn=serve_precomputed_pdb,
        inputs=[seq_input, desc_input],
        outputs=[viewer_html]
    )

demo.launch()