"""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 @dataclass(frozen=True) 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)