"""Public prompt builders for OpticalDNA decoder-style inference.""" from __future__ import annotations from dataclasses import dataclass from enum import Enum from typing import Any, Dict, List class TaskType(str, Enum): """Prompted pretraining / inference task types used by OpticalDNA.""" T1_FULL_OCR = "t1_full_ocr" T2_FULL_OCR_GROUNDING = "t2_full_ocr_grounding" T3_ROI_OCR = "t3_roi_ocr" T4_MASK_COMPLETION = "t4_mask_completion" T5_SUBSEQ_LOCATE = "t5_subseq_locate" T6_CHR_CLASSIFICATION = "t6_chr_classification" class PromptLength(str, Enum): SHORT = "short" MEDIUM = "medium" LONG = "long" @dataclass(frozen=True) class PromptGenerator: """Construct OpticalDNA decoder prompts. This mirrors the prompt families used during OpticalDNA pretraining while providing a small public surface for released-checkpoint inference. """ def build( self, task: TaskType | str, length: PromptLength | str = PromptLength.SHORT, sample: Dict[str, Any] | None = None, ) -> str: task = TaskType(task) length = PromptLength(length) sample = sample or {} if task == TaskType.T1_FULL_OCR: return self._t1(length) if task == TaskType.T2_FULL_OCR_GROUNDING: return self._t2(length) if task == TaskType.T3_ROI_OCR: boxes = sample.get("boxes") if boxes is None: raise KeyError("T3 requires sample={'boxes': [[img_id,x1,y1,x2,y2], ...]}") return self._t3(length, boxes) if task == TaskType.T4_MASK_COMPLETION: boxes = sample.get("boxes") if boxes is None: raise KeyError("T4 requires sample={'boxes': [[img_id,x1,y1,x2,y2], ...]}") return self._t4(length, boxes) if task == TaskType.T5_SUBSEQ_LOCATE: query = sample.get("query") if query is None: raise KeyError("T5 requires sample={'query': 'ACGT...'}") return self._t5(length, query) if task == TaskType.T6_CHR_CLASSIFICATION: return self._t6(length) raise ValueError(f"Unsupported task: {task}") def free_ocr(self, length: PromptLength | str = PromptLength.SHORT) -> str: """Convenience alias for the most common release-time decoder prompt.""" return self._t1(PromptLength(length)) def _t1(self, length: PromptLength) -> str: if length == PromptLength.SHORT: return "Free OCR." if length == PromptLength.MEDIUM: return "Free OCR.\nReturn only DNA sequence." return ( "Free OCR.\n" "Return only the DNA sequence (A/C/G/T/N). " "Keep line breaks if present. No extra words." ) def _t2(self, length: PromptLength) -> str: if length == PromptLength.SHORT: return "<|grounding|>Read all DNA text and locate each line." if length == PromptLength.MEDIUM: return ( "<|grounding|>Read all DNA text and locate each line.\n" "Output per line: <|ref|>SEQ<|/ref|><|det|>[[img_id,x1,y1,x2,y2]]<|/det|>." ) return ( "<|grounding|>" "Read all DNA text and locate each text line or block.\n" "Output one line per region in reading order (top-to-bottom, left-to-right).\n" "For each region, output EXACTLY:\n" "<|ref|>SEQUENCE<|/ref|><|det|>[[img_id,x1,y1,x2,y2]]<|/det|>\n" "Rules:\n" "- det MUST be a list-of-boxes. Even one box must be written as [[...]].\n" "- img_id is 0-based index into the images list. If NUM_IMAGES==1, img_id MUST be 0.\n" "Only output these lines. No extra text." ) def _t3(self, length: PromptLength, boxes: List[List[int]]) -> str: box_str = str(boxes) if length == PromptLength.SHORT: return f"<|grounding|>OCR DNA for boxes (in order): {box_str}" if length == PromptLength.MEDIUM: return ( "<|grounding|>" "OCR DNA text for each box in the SAME order.\n" f"Boxes:\n{box_str}\n" "Output one line per box: <|ref|>SEQ<|/ref|><|det|>[[img_id,x1,y1,x2,y2]]<|/det|>" ) return ( "<|grounding|>" "OCR DNA text for each bounding box below, in the SAME order.\n" f"Boxes:\n{box_str}\n" "Output one line per box using EXACTLY:\n" "<|ref|>SEQUENCE<|/ref|><|det|>[[img_id,x1,y1,x2,y2]]<|/det|>\n" "Rules:\n" "- det MUST be list-of-boxes: [[...]] for a single box.\n" "- img_id is 0-based index into the images list. If NUM_IMAGES==1, img_id MUST be 0.\n" "No extra text." ) def _t4(self, length: PromptLength, boxes: List[List[int]]) -> str: box_str = str(boxes) if length == PromptLength.SHORT: return f"<|grounding|>Predict masked DNA for boxes (in order): {box_str}" if length == PromptLength.MEDIUM: return ( "<|grounding|>" "Predict ORIGINAL DNA for masked boxes in the SAME order.\n" f"Boxes:\n{box_str}\n" "Output: <|ref|>SEQ<|/ref|><|det|>[[img_id,x1,y1,x2,y2]]<|/det|>" ) return ( "<|grounding|>" "The DNA text inside each box is masked/occluded.\n" "Predict the ORIGINAL DNA sequence for each masked region.\n" f"Masked boxes:\n{box_str}\n" "Output in the SAME order as the boxes.\n" "Use only A/C/G/T/N (use N if uncertain).\n" "For each box, output EXACTLY one line:\n" "<|ref|>PREDICTED_SEQUENCE<|/ref|><|det|>[[img_id,x1,y1,x2,y2]]<|/det|>\n" "Rules:\n" "- det MUST be list-of-boxes: [[...]] for a single box.\n" "- img_id is 0-based index into the images list. If NUM_IMAGES==1, img_id MUST be 0.\n" "No extra text." ) def _t5(self, length: PromptLength, query: str) -> str: if length == PromptLength.SHORT: return f"<|grounding|>Locate <|ref|>{query}<|/ref|>." if length == PromptLength.MEDIUM: return ( "<|grounding|>" f"Locate <|ref|>{query}<|/ref|>.\n" "Output exactly one line: <|ref|>QUERY<|/ref|><|det|>[...]<|/det|> (or [] if not found)." ) return ( "<|grounding|>" f"Locate the DNA subsequence <|ref|>{query}<|/ref|>.\n" "Return ALL bounding boxes where it appears.\n" "Output EXACTLY one line:\n" f"<|ref|>{query}<|/ref|><|det|>[[img_id,x1,y1,x2,y2],[img_id,x1,y1,x2,y2]]<|/det|>\n" "If not found, output:\n" f"<|ref|>{query}<|/ref|><|det|>[]<|/det|>\n" "Rules:\n" "- Each box MUST be [img_id,x1,y1,x2,y2].\n" "- img_id is 0-based index into the images list. If NUM_IMAGES==1, img_id MUST be 0.\n" "No extra text." ) def _t6(self, length: PromptLength) -> str: if length == PromptLength.SHORT: return "Which chromosome?" if length == PromptLength.MEDIUM: return "Predict chromosome label (chr1-22, chrX, chrY, or unknown)." return ( "Classify this DNA sequence: which human chromosome does it belong to?\n" "Answer with one label only: chr1-chr22, chrX, chrY, or unknown." )