factorforge-cds 3.0.0__py3-none-any.whl

factorforge/engines/v3/metrics.py

@@ -0,0 +1,145 @@
+"""Metrics for FactorForge v3 (CAI, perplexity, GC%)."""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # pragma: no cover - for type checkers only
+    import torch
+
+
+def _default_codon_table_path() -> Path:
+    return Path(__file__).resolve().parents[4] / "data" / "nbenthamiana_codons.json"
+
+
+@dataclass(frozen=True)
+class CodonUsageTable:
+    codon_to_aa: dict[str, str]
+    codon_weights: dict[str, float]
+    best_codon_for_aa: dict[str, str]
+    source: str | None = None
+
+
+def load_codon_usage_table(path: Path | None = None) -> CodonUsageTable:
+    table_path = path or _default_codon_table_path()
+    raw = json.loads(table_path.read_text(encoding="utf-8"))
+    if not isinstance(raw, dict):
+        raise ValueError("Codon table JSON must be an object")
+
+    codons = raw.get("codons")
+    if not isinstance(codons, dict):
+        raise ValueError("Codon table JSON missing 'codons'")
+
+    codon_to_aa: dict[str, str] = {}
+    codon_freq: dict[str, float] = {}
+    for codon, entry in codons.items():
+        if not isinstance(entry, dict) or not isinstance(codon, str):
+            continue
+        aa = entry.get("aa")
+        freq = entry.get("frequency")
+        if not isinstance(aa, str) or not isinstance(freq, (int, float)):
+            continue
+        codon_to_aa[codon] = aa
+        codon_freq[codon] = float(freq)
+
+    codon_weights = _build_codon_weights(codon_to_aa, codon_freq)
+    best_codon_for_aa = _best_codon_map(codon_to_aa, codon_freq)
+    source = raw.get("source") if isinstance(raw.get("source"), str) else None
+
+    return CodonUsageTable(
+        codon_to_aa=codon_to_aa,
+        codon_weights=codon_weights,
+        best_codon_for_aa=best_codon_for_aa,
+        source=source,
+    )
+
+
+def compute_cai(dna_sequence: str, table: CodonUsageTable) -> float:
+    seq = dna_sequence.upper().replace("U", "T")
+    codon_count = len(seq) // 3
+    if codon_count == 0:
+        return 0.0
+
+    weights: list[float] = []
+    for i in range(codon_count):
+        codon = seq[i * 3 : i * 3 + 3]
+        weight = table.codon_weights.get(codon)
+        if weight is None or weight <= 0:
+            return 0.0
+        weights.append(weight)
+
+    log_sum = sum(math.log(w) for w in weights)
+    return math.exp(log_sum / len(weights))
+
+
+def compute_gc(dna_sequence: str) -> float:
+    seq = dna_sequence.upper()
+    if not seq:
+        return 0.0
+    gc = seq.count("G") + seq.count("C")
+    return (gc / len(seq)) * 100.0
+
+
+def compute_perplexity(
+    logits: "torch.Tensor",
+    labels: "torch.Tensor",
+    ignore_index: int = -100,
+) -> float:
+    torch, functional = _require_torch()
+    vocab_size = int(logits.shape[-1])
+    loss = functional.cross_entropy(
+        logits.view(-1, vocab_size),
+        labels.view(-1),
+        ignore_index=ignore_index,
+    )
+    return float(torch.exp(loss).item())
+
+
+def _build_codon_weights(
+    codon_to_aa: dict[str, str],
+    codon_freq: dict[str, float],
+) -> dict[str, float]:
+    by_aa: dict[str, list[float]] = {}
+    for codon, aa in codon_to_aa.items():
+        if aa == "*":
+            continue
+        by_aa.setdefault(aa, []).append(codon_freq.get(codon, 0.0))
+
+    weights: dict[str, float] = {}
+    for codon, aa in codon_to_aa.items():
+        if aa == "*":
+            continue
+        max_freq = max(by_aa.get(aa, [0.0]))
+        freq = codon_freq.get(codon, 0.0)
+        weights[codon] = freq / max_freq if max_freq > 0 else 0.0
+    return weights
+
+
+def _best_codon_map(
+    codon_to_aa: dict[str, str],
+    codon_freq: dict[str, float],
+) -> dict[str, str]:
+    best: dict[str, tuple[str, float]] = {}
+    for codon, aa in codon_to_aa.items():
+        if aa == "*":
+            continue
+        current = best.get(aa)
+        freq = codon_freq.get(codon, 0.0)
+        if current is None or freq > current[1]:
+            best[aa] = (codon, freq)
+    return {aa: codon for aa, (codon, _) in best.items()}
+
+
+def _require_torch():
+    try:
+        import torch
+        from torch.nn import functional
+    except ImportError as exc:  # pragma: no cover
+        raise ImportError(
+            "ML dependencies not installed. Install with: pip install -e \".[ml]\""
+        ) from exc
+    return torch, functional

factorforge/engines/v3/modeling_bart_decoder.py

@@ -0,0 +1,127 @@
+"""BART decoder skeleton for FactorForge v3."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+try:
+    import torch
+    from torch import nn
+    from transformers import BartConfig
+    from transformers.models.bart.modeling_bart import BartDecoder
+except ImportError as exc:  # pragma: no cover - exercised in ML installs
+    raise ImportError(
+        "ML dependencies not installed. Install with: pip install -e \".[ml]\""
+    ) from exc
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+
+class BartDecoderSkeleton(nn.Module):
+    """Lightweight BART decoder that consumes unified encoder embeddings."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 256,
+        encoder_dim: int = 320,  # ESM2 t6_8M per-token embedding dim
+        decoder_layers: int = 4,
+        decoder_attention_heads: int = 4,
+        ffn_dim: int = 1024,
+        max_position_embeddings: int = 512,
+        dropout: float = 0.1,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+    ) -> None:
+        super().__init__()
+        self.config = BartConfig(
+            vocab_size=vocab_size,
+            d_model=d_model,
+            decoder_layers=decoder_layers,
+            decoder_attention_heads=decoder_attention_heads,
+            decoder_ffn_dim=ffn_dim,
+            max_position_embeddings=max_position_embeddings,
+            encoder_layers=1,
+            encoder_attention_heads=decoder_attention_heads,
+            dropout=dropout,
+            activation_function="gelu",
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=bos_token_id,
+            is_encoder_decoder=True,
+        )
+
+        self.encoder_projection = nn.Linear(encoder_dim, d_model)
+        self.decoder = BartDecoder(self.config)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        self.register_buffer("final_logits_bias", torch.zeros((1, vocab_size)))
+
+    def forward(
+        self,
+        encoder_hidden_states: "Tensor",
+        decoder_input_ids: "Tensor",
+        attention_mask: "Tensor | None" = None,
+        encoder_attention_mask: "Tensor | None" = None,
+    ) -> "Tensor":
+        """Teacher-forcing forward pass returning logits [B, T, vocab_size]."""
+        encoder_states = self.encoder_projection(encoder_hidden_states)
+        outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=False,
+            return_dict=True,
+        )
+        logits = self.lm_head(outputs.last_hidden_state) + self.final_logits_bias.to(
+            outputs.last_hidden_state.device
+        )
+        return logits
+
+    def generate(
+        self,
+        encoder_hidden_states: "Tensor",
+        max_new_tokens: int,
+        bos_token_id: int,
+        eos_token_id: int,
+        pad_token_id: int | None = None,
+    ) -> "Tensor":
+        """Greedy decoding for a few steps (no beam search)."""
+        batch_size = encoder_hidden_states.shape[0]
+        device = encoder_hidden_states.device
+
+        decoder_input_ids = torch.full(
+            (batch_size, 1),
+            bos_token_id,
+            dtype=torch.long,
+            device=device,
+        )
+
+        for _ in range(max_new_tokens):
+            logits = self.forward(
+                encoder_hidden_states=encoder_hidden_states,
+                decoder_input_ids=decoder_input_ids,
+                attention_mask=None,
+                encoder_attention_mask=None,
+            )
+            next_token = torch.argmax(logits[:, -1, :], dim=-1)
+            decoder_input_ids = torch.cat(
+                [decoder_input_ids, next_token.unsqueeze(1)],
+                dim=1,
+            )
+            if torch.all(next_token == eos_token_id):
+                break
+
+        if pad_token_id is not None:
+            decoder_input_ids = decoder_input_ids.clone()
+            for idx in range(decoder_input_ids.shape[0]):
+                row = decoder_input_ids[idx]
+                eos_positions = (row == eos_token_id).nonzero(as_tuple=False)
+                if eos_positions.numel() > 0:
+                    eos_position = int(eos_positions[0])
+                    if eos_position + 1 < row.shape[0]:
+                        row[eos_position + 1 :] = pad_token_id
+        return decoder_input_ids

factorforge/engines/v3/pipeline.py

@@ -0,0 +1,192 @@
+"""FactorForge v3 pipeline scaffolding."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from factorforge.core.interfaces import OptimizationResult, OptimizerEngine
+from factorforge.engines.registry import EngineRegistry
+from factorforge.engines.v2.rules.domesticator import Domesticator
+from factorforge.engines.v2.rules.rule_engine import RuleEngine
+
+from .explain import ExplainabilityInputs, build_fda_report
+from .metrics import CodonUsageTable, compute_cai, compute_gc, load_codon_usage_table
+from .tokenizer import AA_TOKENS, AATokenizer, CodonTokenizer
+
+if TYPE_CHECKING:
+    from .modeling_bart_decoder import BartDecoderSkeleton
+
+
+VALID_AA = set(AA_TOKENS)
+
+
+@dataclass
+class V3Result:
+    sequence: str
+    metrics: dict[str, float]
+    metadata: dict[str, Any] = field(default_factory=dict)
+    report: dict[str, Any] | None = None
+
+
+class V3Pipeline:
+    """Minimal v3 pipeline scaffold (decoder + post-guard hook)."""
+
+    def __init__(
+        self,
+        aa_tokenizer: AATokenizer | None = None,
+        codon_tokenizer: CodonTokenizer | None = None,
+        codon_usage: CodonUsageTable | None = None,
+        decoder: "BartDecoderSkeleton | None" = None,
+        model_id: str = "v3-bart-decoder-skeleton",
+    ) -> None:
+        self.aa_tokenizer = aa_tokenizer or AATokenizer.default()
+        self.codon_tokenizer = codon_tokenizer or CodonTokenizer.default()
+        self.codon_usage = codon_usage or load_codon_usage_table()
+        self.decoder = decoder
+        self.model_id = model_id
+
+    def run(
+        self,
+        sequence: str,
+        encoder_embeddings: Any | None = None,
+        max_new_tokens: int = 24,
+        apply_post_guard: bool = True,
+        seed: int = 0,
+        config: dict[str, Any] | None = None,
+    ) -> V3Result:
+        aa_sequence = _normalize_aa(sequence)
+
+        if encoder_embeddings is not None:
+            if self.decoder is None:
+                raise ValueError("Decoder not configured. Provide a BartDecoderSkeleton instance.")
+            from factorforge.engines.v3.inference.constrained_decoder import (
+                constrained_greedy_decode,
+                validate_candidate_or_fallback,
+            )
+
+            token_ids = constrained_greedy_decode(
+                self.decoder,
+                encoder_embeddings,
+                aa_sequence,
+                self.codon_tokenizer,
+            )
+            dna_sequence = self.codon_tokenizer.decode(token_ids[0].tolist())
+            fallback = validate_candidate_or_fallback(aa_sequence, dna_sequence)
+            dna_sequence = fallback["dna_sequence"]
+        else:
+            dna_sequence = self._fallback_gc_target(aa_sequence)
+
+        post_guard: dict[str, Any] | None = None
+        if apply_post_guard:
+            dna_sequence, post_guard = self._apply_post_guard(dna_sequence)
+
+        metrics = {
+            "cai": compute_cai(dna_sequence, self.codon_usage),
+            "gc_content": compute_gc(dna_sequence),
+        }
+
+        report_inputs = ExplainabilityInputs(
+            aa_sequence=aa_sequence,
+            dna_sequence=dna_sequence,
+            metrics=metrics,
+            model_id=self.model_id,
+            tokenizer_hash=self.codon_tokenizer.mapping_hash(),
+            seed=seed,
+            config=config or {},
+            post_guard=post_guard,
+        )
+        report = build_fda_report(report_inputs)
+
+        metadata = {
+            "model_id": self.model_id,
+            "post_guard": post_guard or {},
+        }
+
+        return V3Result(sequence=dna_sequence, metrics=metrics, metadata=metadata, report=report)
+
+    def _reverse_translate_best(self, aa_sequence: str) -> str:
+        codons: list[str] = []
+        for aa in aa_sequence:
+            codon = self.codon_usage.best_codon_for_aa.get(aa)
+            if codon is None:
+                raise ValueError(f"No codon mapping for amino acid: {aa}")
+            codons.append(codon)
+        return "".join(codons)
+
+    def _fallback_gc_target(self, aa_sequence: str) -> str:
+        from factorforge.engines.v3.inference.v2_adapter import optimize_with_v2
+
+        return optimize_with_v2(aa_sequence, options={"profile": "gc_target"})["dna_sequence"]
+
+    def _apply_post_guard(self, dna_sequence: str) -> tuple[str, dict[str, Any]]:
+        rule_engine = RuleEngine()
+        domesticator = Domesticator()
+
+        # Step 1: PolyA iterative fix (v2 pipeline과 동일한 방식)
+        polya_fix: dict[str, Any] | None = None
+        current_seq = dna_sequence
+        has_polya = any(p in dna_sequence for p in rule_engine.POLYA_PATTERNS)
+        if has_polya:
+            polya_result = rule_engine.fix_polya_iterative(current_seq)
+            if polya_result["success"]:
+                current_seq = polya_result["modified_seq"]
+            polya_fix = polya_result
+
+        # Step 2: Full scan on fixed sequence
+        scan_results = rule_engine.scan_all(current_seq)
+
+        # Step 3: Domestication
+        domestication = domesticator.domesticate(current_seq, standard="golden_gate")
+        domesticated = domestication.get("domesticated_seq", current_seq)
+
+        post_guard = {
+            "scan_results": scan_results,
+            "domestication": domestication,
+            "polya_fix": polya_fix,
+            "edited": domesticated != dna_sequence,
+        }
+        return domesticated, post_guard
+
+
+class V3Optimizer(OptimizerEngine):
+    """Optimizer wrapper for v3 pipeline."""
+
+    name = "v3 BART Decoder"
+    version = "3.0.0"
+
+    def __init__(self, pipeline: V3Pipeline | None = None) -> None:
+        self.pipeline = pipeline or V3Pipeline()
+
+    def optimize(
+        self,
+        sequence: str,
+        profile: str | None = None,
+        **kwargs: Any,
+    ) -> OptimizationResult:
+        result = self.pipeline.run(sequence, **kwargs)
+        return OptimizationResult(
+            sequence=result.sequence,
+            metrics=result.metrics,
+            metadata=result.metadata,
+        )
+
+    def validate(self, sequence: str) -> bool:
+        try:
+            _normalize_aa(sequence)
+            return True
+        except ValueError:
+            return False
+
+    def get_supported_profiles(self) -> list[str]:
+        return ["default"]
+
+
+def _normalize_aa(sequence: str) -> str:
+    seq = sequence.strip().replace(" ", "").replace("\n", "").upper()
+    if not seq:
+        raise ValueError("Protein sequence is empty.")
+    invalid = {aa for aa in seq if aa not in VALID_AA}
+    if invalid:
+        raise ValueError(f"Invalid amino acids found: {''.join(sorted(invalid))}")
+    return seq

factorforge/engines/v3/synonym_mask.py

@@ -0,0 +1,61 @@
+"""Synonymous codon masks for v3 training and constrained decoding."""
+
+from __future__ import annotations
+
+from typing import Mapping
+
+import torch
+
+from factorforge.ml.metrics import STANDARD_GENETIC_CODE
+
+
+AA_TO_CODONS: dict[str, tuple[str, ...]] = {}
+for _codon, _aa in STANDARD_GENETIC_CODE.items():
+    AA_TO_CODONS.setdefault(_aa, tuple())
+for _aa in list(AA_TO_CODONS):
+    AA_TO_CODONS[_aa] = tuple(
+        codon for codon, codon_aa in STANDARD_GENETIC_CODE.items() if codon_aa == _aa
+    )
+
+
+def normalize_protein_sequence(protein_sequence: str) -> str:
+    """Normalize a protein sequence for codon-mask construction."""
+    protein = "".join(protein_sequence.upper().split())
+    if protein.endswith("*"):
+        protein = protein[:-1]
+    if not protein:
+        raise ValueError("protein_sequence must not be empty")
+    invalid = [aa for aa in protein if aa not in AA_TO_CODONS or aa == "*"]
+    if invalid:
+        raise ValueError(f"No synonymous codons for amino acid(s): {''.join(sorted(set(invalid)))}")
+    return protein
+
+
+def synonymous_codons_for_aa(amino_acid: str) -> tuple[str, ...]:
+    """Return synonymous non-stop codons for one standard amino acid."""
+    aa = amino_acid.upper()
+    codons = AA_TO_CODONS.get(aa, tuple())
+    if not codons or aa == "*":
+        raise ValueError(f"No synonymous codons for amino acid: {amino_acid}")
+    return codons
+
+
+def build_synonym_token_mask(
+    protein_sequence: str,
+    token_to_id: Mapping[str, int],
+    *,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    """Build a boolean mask [protein_length, vocab_size] for synonymous codons only."""
+    protein = normalize_protein_sequence(protein_sequence)
+    vocab_size = max(token_to_id.values()) + 1
+    mask = torch.zeros((len(protein), vocab_size), dtype=torch.bool, device=device)
+    for index, aa in enumerate(protein):
+        for codon in synonymous_codons_for_aa(aa):
+            token_id = token_to_id.get(codon)
+            if token_id is not None:
+                mask[index, token_id] = True
+        if not bool(mask[index].any()):
+            raise ValueError(f"No tokenizer codons available for amino acid: {aa}")
+    return mask
+