Spaces:
Running on Zero
Running on Zero
| """Gradio API adapter for the ASU CodonTransformer wiki.""" | |
| from __future__ import annotations | |
| import importlib | |
| import os | |
| from pathlib import Path | |
| import gradio as gr | |
| import spaces | |
| from huggingface_hub import hf_hub_download | |
| from api.huggingface import register_adp1_organism | |
| from api.service import ( | |
| MAX_FASTA_BYTES, | |
| CodonTransformerService, | |
| _ensure_codontransformer_import, | |
| parse_dna_text, | |
| ) | |
| _ensure_codontransformer_import() | |
| codon_utils = importlib.import_module("CodonTransformer.CodonUtils") | |
| register_adp1_organism(codon_utils) | |
| model_path = os.environ.get("CODONTRANSFORMER_MODEL_PATH") | |
| model_repo_id = os.environ.get("CODONTRANSFORMER_MODEL_REPO_ID", "").strip() | |
| if not model_path and model_repo_id: | |
| model_path = hf_hub_download( | |
| repo_id=model_repo_id, | |
| filename=os.environ.get( | |
| "CODONTRANSFORMER_MODEL_FILENAME", | |
| "a_baylyi_adp1/a_baylyi_adp1.ckpt", | |
| ), | |
| ) | |
| os.environ["CODONTRANSFORMER_MODEL_PATH"] = model_path | |
| elif not model_path: | |
| bundled_checkpoint = ( | |
| Path(__file__).parent | |
| / "checkpoints" | |
| / "a_baylyi_adp1" | |
| / "a_baylyi_adp1.ckpt" | |
| ) | |
| if bundled_checkpoint.is_file(): | |
| model_path = str(bundled_checkpoint) | |
| os.environ["CODONTRANSFORMER_MODEL_PATH"] = model_path | |
| predictor = CodonTransformerService() | |
| def list_organisms() -> list[dict[str, int | str]]: | |
| return predictor.supported_organisms | |
| def predict(organism: str, sequence_text: str) -> dict[str, list[dict[str, str]]]: | |
| if len(sequence_text.encode("utf-8")) > MAX_FASTA_BYTES: | |
| raise ValueError("The FASTA input must be 25 MB or smaller.") | |
| records = parse_dna_text(sequence_text) | |
| return {"results": predictor.predict(records, organism)} | |
| with gr.Blocks(title="ASU CodonTransformer") as demo: | |
| gr.Markdown( | |
| "# ASU CodonTransformer\n" | |
| "Optimize coding DNA for an organism supported by the loaded model." | |
| ) | |
| organism_input = gr.Textbox(label="Target organism") | |
| sequence_input = gr.Textbox( | |
| label="DNA sequence or FASTA", | |
| lines=8, | |
| placeholder=">gene1\nATGAAACCCGGGTAA", | |
| ) | |
| organisms_output = gr.JSON(visible=False) | |
| gr.Button("List supported organisms").click( | |
| fn=list_organisms, | |
| inputs=[], | |
| outputs=organisms_output, | |
| api_name="organisms", | |
| ) | |
| prediction_output = gr.JSON(label="Prediction") | |
| gr.Button("Optimize DNA").click( | |
| fn=predict, | |
| inputs=[organism_input, sequence_input], | |
| outputs=prediction_output, | |
| api_name="predict", | |
| ) | |
| demo.queue() | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.environ.get("PORT", "7860")), | |
| show_error=True, | |
| ) | |