app.py

# --- Imports ---
import gradio as gr
import spaces
from transformers import pipeline

# --- Load Model ---
pipe = pipeline(model="InstaDeepAI/ChatNT", trust_remote_code=True)

# --- Logs ---
log_file = "logs.txt"


def log_message(message: str):
    with open(log_file, "a") as log:
        log.write(f"{message}\n")


# --- Utilities ---
def read_dna_sequence(dna_text, fasta_file):
    """
    This function reads the DNA sequence, either from the text field or from the FASTA file,
    and it checks the format.
    It might return an error if the format of the FASTA file is incorrect.
    
    Returns:
        dna_sequence: str
        warning: str if any
        error: str if any
    """
    dna_sequence = ""
    warning = ""
    error = ""

    # Pasted text
    if dna_text and dna_text.strip():
        dna_sequence = dna_text.strip().replace("\n", "")

    # Text overrides uploaded FASTA
    if fasta_file is not None:
        if dna_sequence:
            warning = "⚠️ Warning: Both pasted DNA and FASTA file provided. Using the pasted DNA only."

        else:
            try:
                with open(fasta_file.name, "r") as f:
                    content = f.read()
    
                if not content.startswith(">"):
                    error = "❌ Invalid FASTA: must start with '>' header line."
                    return "", warning, error
    
                sequence = ""
                for line in content.splitlines():
                    if not line or line.startswith(">"):
                        continue
                    sequence += line.strip()
    
                dna_sequence = sequence
    
            except Exception:
                error = "❌ Could not read the FASTA file."

    if dna_sequence and not dna_sequence.isupper():
        dna_sequence = dna_sequence.upper()
        warning += "\n⚠️ Note: DNA sequence was converted to uppercase (lowercase nucleotides are not supported)."

    if error == "":
        # If there was an error before, no need to check the content of the sequence. Otherwise, check that only ACGTN are used
        for nucleotide in dna_sequence:
            if nucleotide not in ["A", "C", "G", "T", "N"]:
                error = "❌ The nucleotides of the DNA sequence should be either 'A', 'T', 'G', 'C' or 'N'."

    return dna_sequence, warning.strip(), error


def validate_inputs(dna_sequence: str, custom_question: str):
    """
    This function is used to

    Args:
        dna_sequence (str): DNA sequence used by the user.
            Note: It might be empty ("")

        custom_question (str): Question of the user (in english).
        
    Returns:
        valid: bool
            Whether the input is valid (correct number of <DNA> tokens, etc)
        error: str
            Error message if necessary
    """
    placeholders = custom_question.count("<DNA>")

    if not custom_question.strip():
        return False, "❌ Please provide a question."

    if dna_sequence and placeholders == 0:
        log_message("Error: DNA sequence provided but no <DNA> token.")
        return False, "❌ Your question must contain the <DNA> token if you provide a DNA sequence."

    if not dna_sequence and placeholders == 1:
        log_message("Error: <DNA> token but no sequence.")
        return False, "❌ You must provide a DNA sequence if you use the <DNA> token."

    if placeholders > 1:
        return False, "❌ Only one <DNA> token is allowed."

    return True, ""


# --- Main Inference ---
@spaces.GPU
def run_chatnt(dna_text, fasta_file, custom_question):
    feedback_msgs = []

    dna_sequence, warning, fasta_error = read_dna_sequence(dna_text, fasta_file)

    if fasta_error:
        return "", fasta_error

    is_valid, validation_error = validate_inputs(dna_sequence, custom_question)
    if not is_valid:
        return "", validation_error

    final_prompt = custom_question

    inputs = {
        "english_sequence": final_prompt,
        "dna_sequences": [dna_sequence] if dna_sequence else []
    }

    output = pipe(inputs=inputs)
    result = output

    if warning:
        feedback_msgs.append(warning)

    if len(feedback_msgs) == 0:
        feedback_msgs.append("✅ Execution was succesful.")
    return result, "\n".join(feedback_msgs)


# --- Gradio Interface ---
css = """
.gradio-container { font-family: sans-serif; }
.gr-button { color: white; border-color: black; background: black; }
footer { display: none !important; }
"""

example_dna = "AGGCGGTGGCAACAGCAACGATAGTTGCATCAGCGCTAACTGCTGCTGCTGCTGCGTGTTCTGCCGGAGGCTGCCGGGAGGGTACAGTAATGACTGTAGGGAGCGCACAGCCGCTGCGGCGGCAGGTTCCCTGGATCAGGAAATTAGGACAGGCACCCGGGATTGGAGGCGAGGGCGGCGCGCGAGCCAACCAGCGGTCGCCCCGGGCCCCCGTGAGCCCCAGGCGAACGCGCGGAGCAGCCGGGCCCCTAGAGCGGTCCGGGTGGCGGCGGCGGCAGCCACGACCACGACCGCGGTCAC"
example_question = "Is there a promoter element in human or mouse cells present in the nucleotide sequence <DNA>?"

with gr.Blocks(css=css) as demo:
    gr.Markdown(
        """
        # 🧬 ChatNT: A Multimodal Conversational Agent for DNA, RNA and Protein Tasks
        [ChatNT](https://www.nature.com/articles/s42256-025-01047-1) is the first multimodal conversational agent designed with a deep understanding of biological sequences (DNA, RNA, proteins). It enables users — even those with no coding background — to interact with biological data through natural language and it generalizes across multiple biological tasks and modalities.
        
        This Hugging Face Space is powered by a [ZeroGPU](https://huggingface.co/docs/hub/en/spaces-zerogpu), which is free but **limited to 5 minutes per day per user**.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            dna_text = gr.Textbox(
                label="Paste your DNA sequence",
                placeholder="ATGCATGC...",
                lines=4
            )

            fasta_file = gr.File(
                label="Or upload your FASTA file",
                file_types=[".fasta", ".fa", ".txt"],
                height=50
            )

            custom_question = gr.Textbox(
                label="English Question",
                placeholder="Does this sequence <DNA> contain a donor splice site?"
            )

            use_example = gr.Button("Use Example")
            submit_btn = gr.Button("Run Query", variant="primary")

        with gr.Column(scale=1):
            output = gr.Textbox(
                label="Model Answer",
                lines=12
            )
            error_box = gr.Textbox(
                label="Execution Feedback",
                lines=4
            )

    submit_btn.click(
        run_chatnt,
        inputs=[dna_text, fasta_file, custom_question],
        outputs=[output, error_box],
    )

    use_example.click(
        lambda: (example_dna, None, example_question),
        inputs=[],
        outputs=[dna_text, fasta_file, custom_question]
    )

    with gr.Row():
        with gr.Column(scale=3):
            gr.Markdown(
                """
You must use **exactly one `<DNA>` token** if you want the model to see your sequence. 
It is also possible to use the model without any DNA sequence (in this case, the `<DNA>` token must not be present in the question).
You can either paste a sequence or upload a FASTA file.  

---

### Limitations and Disclaimer
ChatNT can only handle questions related to the 27 tasks it has been trained on, including the same format of DNA sequences. ChatNT is not a clinical or diagnostic tool. It can produce incorrect or “hallucinated” answers, particularly on out‑of‑distribution inputs, and its numeric predictions may suffer digit‑level errors. Confidence estimates require post‑hoc calibration. Users should always validate critical outputs against experiments or specialized bioinformatics pipelines.

---

### How to query the model ?
ChatNT works like an advanced assistant for analyzing DNA sequences — but it has some important rules.

#### ✅ What works well
ChatNT is good at answering biological questions about the meaning or function of the DNA sequence.
Examples of good queries:
- "Does this sequence `<DNA>` contain a donor splice site?"
- "Is it possible for you to identify whether there's a substantial presence of H3 histone protein occupancy in the nucleotide sequence `<DNA>` in yeast?"
- "Determine the degradation rate of the mouse RNA sequence `<DNA>` within the -5 to 5 range."

#### ❌ What does not work properly
ChatNT does not have direct access to the raw nucleotides of the DNA. Instead, it works with an internal representation of the sequence. This means it cannot answer questions about the exact nucleotides or modify the sequence for you.
Examples of bad queries:
- "What is the length of this sequence `<DNA>`?"
- "Replace the all 'A' nucleotides by 'G' in this sequence `<DNA>`".

For more examples, you can refer to the [training dataset](https://huggingface.co/datasets/InstaDeepAI/ChatNT_training_data).

---

                """
            )
        with gr.Column(scale=1):
            gr.Image(
                "https://media.springernature.com/w440/springer-static/cover-hires/journal/42256/7/6",
                label=None,
                show_label=False
            )

    gr.Markdown("""
### 📚 Citation

If you use **ChatNT**, please cite:  

```bibtex
@article{deAlmeida2025,
  title  = {A multimodal conversational agent for DNA, RNA and protein tasks},
  author = {de Almeida, Bernardo P. and Richard, Guillaume and Dalla-Torre, Hugo and Blum, Christopher and Hexemer, Lorenz and Pandey, Priyanka and Laurent, Stefan and Rajesh, Chandana and Lopez, Marie and Laterre, Alexandre and Lang, Maren and Şahin, Uğur and Beguir, Karim and Pierrot, Thomas},
  journal = {Nature Machine Intelligence},
  year   = {2025},
  volume = {7},
  number = {6},
  pages  = {928--941},
  doi    = {10.1038/s42256-025-01047-1},
  url    = {https://doi.org/10.1038/s42256-025-01047-1},
  issn   = {2522-5839}
}
```
""",
        show_copy_button=True
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(debug=True, show_error=True)