Spaces:
Running on Zero
Running on Zero
| """Input validation and CodonTransformer model inference.""" | |
| from __future__ import annotations | |
| import io | |
| import importlib | |
| import os | |
| import re | |
| import sys | |
| from dataclasses import dataclass | |
| from threading import Lock | |
| from typing import BinaryIO | |
| from Bio import SeqIO | |
| MAX_FASTA_BYTES = 25 * 1024 * 1024 | |
| MAX_RECORDS = 20 | |
| MAX_CODONS = 2046 | |
| DNA_PATTERN = re.compile(r"^[ACGT]+$", re.IGNORECASE) | |
| FASTA_EXTENSIONS = {".fasta", ".fa", ".fna"} | |
| def _ensure_codontransformer_import() -> None: | |
| """Handle package installers that normalize its import directory to lowercase.""" | |
| try: | |
| importlib.import_module("CodonTransformer") | |
| except ModuleNotFoundError as error: | |
| if error.name != "CodonTransformer": | |
| raise | |
| package = importlib.import_module("codontransformer") | |
| sys.modules["CodonTransformer"] = package | |
| class SequenceInput: | |
| header: str | |
| dna: str | |
| def parse_dna_text(text: str) -> list[SequenceInput]: | |
| """Accept one pasted sequence, or a multi-record FASTA document.""" | |
| cleaned_text = text.lstrip("\ufeff") | |
| if not cleaned_text.strip(): | |
| raise ValueError("Enter a DNA sequence or upload a FASTA file.") | |
| has_fasta_header = any( | |
| line.lstrip().startswith(">") for line in cleaned_text.splitlines() | |
| ) | |
| if has_fasta_header: | |
| if not cleaned_text.lstrip().startswith(">"): | |
| raise ValueError("FASTA records must begin with a header starting with >.") | |
| records = list(SeqIO.parse(io.StringIO(cleaned_text), "fasta")) | |
| if not records: | |
| raise ValueError("No FASTA records were found.") | |
| parsed = [ | |
| SequenceInput(header=record.description, dna=str(record.seq)) | |
| for record in records | |
| ] | |
| else: | |
| parsed = [SequenceInput(header="Pasted sequence", dna=cleaned_text)] | |
| if len(parsed) > MAX_RECORDS: | |
| raise ValueError(f"Upload no more than {MAX_RECORDS} sequences at a time.") | |
| for index, record in enumerate(parsed, start=1): | |
| if not record.header.strip(): | |
| raise ValueError(f"Sequence {index} needs a non-empty FASTA header.") | |
| dna = re.sub(r"\s+", "", record.dna).upper() | |
| if not dna: | |
| raise ValueError(f"Sequence {index} is empty.") | |
| if not DNA_PATTERN.fullmatch(dna): | |
| raise ValueError( | |
| f"Sequence {index} must contain only A, C, G, and T (no N or other symbols)." | |
| ) | |
| if len(dna) % 3: | |
| raise ValueError( | |
| f"Sequence {index} contains {len(dna)} bases; its length must be divisible by three." | |
| ) | |
| if len(dna) // 3 > MAX_CODONS: | |
| raise ValueError( | |
| f"Sequence {index} exceeds the {MAX_CODONS}-codon model limit." | |
| ) | |
| parsed[index - 1] = SequenceInput(header=record.header, dna=dna) | |
| return parsed | |
| def read_fasta_upload(upload: BinaryIO, filename: str) -> list[SequenceInput]: | |
| """Read and validate an uploaded FASTA file.""" | |
| suffix = os.path.splitext(filename)[1].lower() | |
| if suffix not in FASTA_EXTENSIONS: | |
| raise ValueError("Upload a .fasta, .fa, or .fna file.") | |
| contents = upload.read(MAX_FASTA_BYTES + 1) | |
| if len(contents) > MAX_FASTA_BYTES: | |
| raise ValueError("The FASTA file must be 25 MB or smaller.") | |
| try: | |
| fasta_text = contents.decode("utf-8-sig") | |
| except UnicodeDecodeError as error: | |
| raise ValueError("The FASTA file must be UTF-8 text.") from error | |
| if not fasta_text.lstrip().startswith(">"): | |
| raise ValueError("FASTA records must begin with a header starting with >.") | |
| return parse_dna_text(fasta_text) | |
| class CodonTransformerService: | |
| """Load the model once and predict organism-optimized coding DNA.""" | |
| def __init__(self) -> None: | |
| import torch | |
| _ensure_codontransformer_import() | |
| from CodonTransformer.CodonData import get_amino_acid_sequence, get_codon_table | |
| from CodonTransformer.CodonPrediction import ( | |
| load_model, | |
| load_tokenizer, | |
| predict_dna_sequence, | |
| ) | |
| from CodonTransformer.CodonUtils import ( | |
| MAX_AMINO_ACIDS, | |
| ORGANISM2ID, | |
| STOP_SYMBOL, | |
| ) | |
| self.torch = torch | |
| self.get_amino_acid_sequence = get_amino_acid_sequence | |
| self.get_codon_table = get_codon_table | |
| self.predict_dna_sequence = predict_dna_sequence | |
| self.organisms = ORGANISM2ID | |
| self.stop_symbol = STOP_SYMBOL | |
| self.max_amino_acids = MAX_AMINO_ACIDS | |
| self._prediction_lock = Lock() | |
| if torch.cuda.is_available(): | |
| self.device = torch.device("cuda") | |
| elif torch.backends.mps.is_available(): | |
| self.device = torch.device("mps") | |
| else: | |
| self.device = torch.device("cpu") | |
| model_path = os.environ.get("CODONTRANSFORMER_MODEL_PATH") or None | |
| self.tokenizer = load_tokenizer() | |
| self.model = load_model(model_path=model_path, device=self.device) | |
| self.supported_organisms = self._supported_organism_options() | |
| def _supported_organism_options(self) -> list[dict[str, int | str]]: | |
| embedding_count = ( | |
| self.model.bert.embeddings.token_type_embeddings.num_embeddings | |
| ) | |
| return [ | |
| {"id": organism_id, "name": name} | |
| for name, organism_id in sorted( | |
| self.organisms.items(), key=lambda organism: organism[1] | |
| ) | |
| if organism_id < embedding_count | |
| ] | |
| def predict( | |
| self, records: list[SequenceInput], organism: str | |
| ) -> list[dict[str, str]]: | |
| if organism not in self.organisms: | |
| raise ValueError("Select an organism supported by CodonTransformer.") | |
| if not any(option["name"] == organism for option in self.supported_organisms): | |
| raise ValueError( | |
| f"The configured model does not have a trained embedding for {organism}." | |
| ) | |
| results = [] | |
| codon_table = self.get_codon_table(organism) | |
| for record in records: | |
| protein = self.get_amino_acid_sequence( | |
| record.dna, stop_symbol=self.stop_symbol, codon_table=codon_table | |
| ) | |
| if self.stop_symbol in protein[:-1]: | |
| raise ValueError( | |
| f"{record.header}: DNA contains an internal stop codon." | |
| ) | |
| if not protein.endswith(self.stop_symbol): | |
| protein += self.stop_symbol | |
| if len(protein) > self.max_amino_acids: | |
| raise ValueError( | |
| f"{record.header}: translated protein exceeds the model length limit." | |
| ) | |
| with self._prediction_lock: | |
| prediction = self.predict_dna_sequence( | |
| protein=protein, | |
| organism=organism, | |
| device=self.device, | |
| tokenizer=self.tokenizer, | |
| model=self.model, | |
| deterministic=True, | |
| match_protein=True, | |
| ) | |
| simplified_codons = self._simplify_codons( | |
| prediction.predicted_dna, protein, codon_table | |
| ) | |
| results.append( | |
| { | |
| "header": record.header, | |
| "protein": protein, | |
| "optimized_dna": prediction.predicted_dna, | |
| "codon_sequence": simplified_codons, | |
| } | |
| ) | |
| return results | |
| def _simplify_codons(self, dna: str, protein: str, codon_table: int) -> str: | |
| translated = self.get_amino_acid_sequence( | |
| dna, stop_symbol=self.stop_symbol, codon_table=codon_table | |
| ) | |
| if translated != protein: | |
| raise RuntimeError( | |
| "CodonTransformer returned DNA that does not match the input protein." | |
| ) | |
| tokens = [ | |
| f"{amino_acid}_{dna[index * 3:index * 3 + 3]}" | |
| for index, amino_acid in enumerate(translated) | |
| ] | |
| return " ".join(tokens) | |