cradle.sdk.types.tasks.v2

   1from __future__ import annotations
   2
   3from enum import StrEnum
   4from typing import Annotated, Any, ClassVar, Literal
   5
   6from pydantic import BaseModel, Field
   7
   8from .common import (
   9    ArtifactParam,
  10    CustomPredictorParam,
  11    ParametersBase,
  12    ResultBase,
  13    TableResult,
  14    TableSourceLiteral,
  15    TableSourceQuery,
  16    TableSourceSnapshot,
  17    TableSourceTable,
  18    TableSourceTaskResult,
  19)
  20
  21
  22class OptimizationDirection(StrEnum):
  23    MAXIMIZE = "MAXIMIZE"
  24    MINIMIZE = "MINIMIZE"
  25
  26
  27class ScaleType(StrEnum):
  28    MULTIPLICATIVE = "MULTIPLICATIVE"
  29    ADDITIVE = "ADDITIVE"
  30    RANK = "RANK"
  31
  32
  33class AllowedMutation(BaseModel):
  34    in_range: RangeAnnotation = Field(
  35        description="The annotation expression defining the positions at which the amino acids listed in `amino_acids` can be used."
  36    )
  37    amino_acids: list[Annotated[str, Field(description="A single amino acid.")]]
  38
  39
  40class AnalyzeDataParameters(ParametersBase):
  41    """Parameters for the analyze.data task."""
  42
  43    task_type: ClassVar[str] = "analyze.data/v2"
  44
  45    dataset: (
  46        list
  47        | Annotated[
  48            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
  49            Field(discriminator="kind"),
  50        ]
  51    ) = Field(description="The data to analyze before training.")
  52    reference: MonomerAssembly | VhVlAssembly = Field(
  53        description="The reference assembly used for computing mutations."
  54    )
  55    assays: list[Assay] = Field(description="The assays present in the dataset.")
  56
  57
  58class AnalyzeDataResult(ResultBase):
  59    """Result of the analyze.data task."""
  60
  61    report: ArtifactParam = Field(description="The analysis report.")
  62
  63
  64class Assay(BaseModel):
  65    """Metadata of an assay."""
  66
  67    assay_id: str = Field(description="ID of the assay", min_length=1, max_length=124)
  68    name: str = Field(description="Human readable name of the assay", min_length=1, max_length=100)
  69    scale_type: ScaleType = Field(
  70        description="The scale type defines how the assay behaves.\n\n* **Additive**: Assay values are comparable in magnitude over batches, a delta of eg 5 in one batch is equivalent to 5 in another batch.\n* **Multiplicative**: Assay values are comparable over batches by applying a multiplication factor, eg fold improvement: `score_variantB_batch1 = factor(typically starting sequence score) * score_variantB_batch`.\n* **Rank**: Assay values are not comparable over batches, we can only assume the ranking within a batch is correct."
  71    )
  72    unit: str | None = Field(
  73        default=None, description="Unit of the assay - used for display purposes only", min_length=1, max_length=100
  74    )
  75
  76
  77class AssemblyFeature(BaseModel):
  78    ranges: dict[str, list[tuple[int, int]]] = Field(description="The ranges of the feature, keyed by polymer name")
  79    metadata: dict[str, str | int | float | bool] | None = Field(
  80        default=None, description="Optional metadata for the feature"
  81    )
  82
  83
  84class CombinatorialGenerator(BaseModel):
  85    """Combinatorial assembly generation from explicit mutation.
  86
  87    Systematically generates assemblies by combining allowed mutations up to a
  88    specified maximum number.
  89    """
  90
  91    type: Literal["COMBINATORIAL"] = Field(default="COMBINATORIAL")
  92    template: MonomerAssembly | VhVlAssembly = Field(
  93        description="The assembly that serves as template for generation, including any annotations to be used in `requirements`."
  94    )
  95    max_mutations: int = Field(
  96        description="Maximum number of mutations to combine in generation. Lower values will result in fewer total candidate sequences evaluated for this generator."
  97    )
  98    requirements: list[Requirement] = Field(
  99        default_factory=list,
 100        description="A list of requirements to satisfy. All requirements must be met for a valid assembly. See the `Requirement` documentation for details on how requirements are evaluated.",
 101    )
 102    allowed_mutations: list[AllowedMutation] = Field(
 103        description="The complete set of mutations that can be used in generation."
 104    )
 105    alias: str = Field(
 106        default="combinatorial_generator",
 107        description="Custom alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. For example, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
 108    )
 109
 110
 111class Constraint(BaseModel):
 112    """A constraint on an assay.
 113
 114    Multiple constraints can be specified for the same assay, allowing you to define ranges or
 115    multiple boundaries. All constraints must be satisfied for an assembly to be considered valid.
 116    Constraints with the same direction are also supported (e.g., better than reference A, B, and C),
 117    though this is less common.
 118
 119    Example use case: While optimizing for thermostability as the primary objective, you can constrain
 120    binding Kd to stay within an acceptable range by creating two relative constraints on the same
 121    binding assay. Since lower Kd means stronger binding:
 122    - MINIMIZE direction relative to a weak binder reference (high Kd) ensures Kd stays below the weak binder (binding doesn't get too weak)
 123    - MAXIMIZE direction relative to a strong binder reference (low Kd) ensures Kd stays above the strong binder (binding doesn't get too strong)
 124    """
 125
 126    measure: Measure = Field(description="The measure to constrain.")
 127    threshold: float | RelativeTo = Field(
 128        description="The threshold for the constraint, either an absolute value or relative to a reference.\n\n        Note: Relative thresholds (using `RelativeTo`) are strongly preferred over absolute thresholds.\n        Models are trained using contrastive loss, meaning they learn that assembly A > assembly B for\n        a given assay, rather than learning absolute assay values. Consequently, constraining assemblies\n        to be greater than a reference assembly is more likely to produce accurate results than setting\n        an absolute threshold value."
 129    )
 130
 131
 132class DiversifyParameters(ParametersBase):
 133    """Parameters for the diversify task.
 134
 135    An important detail is that the `LikelihoodBasedRanker` expects all assemblies to be of the same shape and length as its reference assembly. See the
 136    `LikelihoodBasedRanker` docs for further details. In particular, `diversify/v2` is not compatible with generators that produce assemblies of varying
 137    shapes/sizes different from the reference assembly.
 138    """
 139
 140    task_type: ClassVar[str] = "diversify/v2"
 141
 142    generator: Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator = Field(
 143        description="The generator configuration or configurations that define requirements and strategies for generating candidate assemblies.",
 144        discriminator="type",
 145    )
 146    ranker: LikelihoodBasedRanker = Field(
 147        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
 148    )
 149    num_assemblies: int = Field(
 150        description="The target number of `selected_assemblies` to output. In rare cases, the actual number of output assemblies may be slightly lower."
 151    )
 152
 153
 154class DiversifyResult(ResultBase):
 155    """Result of the diversify task."""
 156
 157    selected_assemblies: TableResult = Field(
 158        description="Table containing the selected assemblies output by `diversify`."
 159    )
 160    report: ArtifactParam | None = Field(description="A report on the result of the task.")
 161
 162
 163class EngineerParameters(ParametersBase):
 164    """Parameters for the engineer task."""
 165
 166    task_type: ClassVar[str] = "engineer/v2"
 167
 168    generator: Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator = Field(
 169        description="The generator configuration or configurations that define requirements and strategies for generating candidate assemblies.",
 170        discriminator="type",
 171    )
 172    ranker: ScoreBasedRanker = Field(
 173        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
 174    )
 175    num_assemblies: int = Field(
 176        description="The target number of selected assemblies to output. In rare cases of very difficult `constraints` this task may produce a smaller number."
 177    )
 178    filterers: list[ScoreBasedFilterer] = Field(
 179        default_factory=list,
 180        description="If specified, only generated assemblies accepted by all filterers are considered for inclusion in `selected_assemblies`.\n\nFilterers are applied sequentially in the order they are specified. This implies that filterers involving more expensive computations should be placed later\nin the list and, conversely, filterers that are expected to filter out a large portion of the candidate assemblies should be placed earlier in the list.\n        ",
 181    )
 182
 183
 184class EngineerResult(ResultBase):
 185    """Result of the engineer task."""
 186
 187    selected_assemblies: TableResult = Field(
 188        description="Table containing the selected assemblies output by `engineer`."
 189    )
 190    report: ArtifactParam | None = Field(description="A report on the result of the task.")
 191
 192
 193class Explicit(BaseModel):
 194    """A pre-defined table of assemblies to use directly for generation without modification.
 195
 196    Assemblies will be evaluated in the order provided.
 197    """
 198
 199    type: Literal["EXPLICIT"] = Field(default="EXPLICIT")
 200    source: MonomerSource | VhVlSource
 201    alias: str = Field(
 202        default="explicit_generator",
 203        description="Custom alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. For example, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
 204    )
 205
 206
 207class LikelihoodBasedRanker(BaseModel):
 208    """Configuration for ranking and selecting generated assemblies based on likelihoods.
 209
 210    Crucially, this ranking strategy relies on a reference assembly as a basis of comparison (see `reference_assembly` in `sampler_config`) and it
 211    will fail if it ever encounters an assembly not of the same shape/size as this reference. This implies that, unlike `ScoreBasedRanker`, this
 212    ranking strategy only works for batches of assemblies whose polymers are all of the same length.
 213    """
 214
 215    type: Literal["LIKELIHOOD_BASED"] = Field(default="LIKELIHOOD_BASED")
 216    sampler_config: MonomerSamplerConfig | VhVlSamplerConfig = Field(
 217        description="The trained sampler artifact used to compute assembly likelihoods, along with any optional control assemblies to include in the final output for experimental validation.",
 218        discriminator="kind",
 219    )
 220    hit_redundancy: int | None = Field(
 221        default=None,
 222        description='The maximum number of similar variants within the proposed set.\n\nWe say two variants are "similar" when their predicted performances are highly correlated.\nAt lower values this will mean on average fewer hits, but greater diversity within the set,\nwhile at maximum value (`hit_redundancy` equals the number of assemblies to select) the selection is purely greedy, with\nno requirement on diversity within the set. If set to None, it will be automatically set\nto a default value of `1 + num_assemblies//20`.',
 223    )
 224
 225
 226class Measure(BaseModel):
 227    """A measure to optimize or constrain."""
 228
 229    assay_id: str = Field(description="The ID of the assay.")
 230    direction: OptimizationDirection = Field(description="The direction of optimization for the assay.")
 231
 232
 233class ModelBasedGenerator(BaseModel):
 234    """Uses a trained sampler model to generate assemblies."""
 235
 236    type: Literal["MODEL_BASED"] = Field(default="MODEL_BASED")
 237    template: MonomerSamplingTemplate | VhVlSamplingTemplate = Field(
 238        description="The template sequence to use for generation.", discriminator="kind"
 239    )
 240    requirements: list[Requirement] = Field(
 241        default_factory=list,
 242        description="A list of requirements to satisfy. All requirements must be met for a valid assembly. See the `Requirement` documentation for details on how requirements are evaluated.",
 243    )
 244    discourage_mutations: RangeAnnotation | None = Field(
 245        default=None, description="Expression defining regions where mutations should be discouraged."
 246    )
 247    min_mutations: int = Field(description="Minimum number of mutations from the template assembly.")
 248    max_mutations: int = Field(description="Maximum number of mutations from the template assembly.")
 249    alias: str = Field(
 250        default="model_based_generator",
 251        description="Custom alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. For example, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
 252    )
 253
 254
 255class MonomerAssembly(BaseModel):
 256    """A single polymer chain assembly."""
 257
 258    monomer: str = Field(description="The monomer sequence")
 259
 260
 261class MonomerData(BaseModel):
 262    kind: Literal["MONOMER"] = Field(default="MONOMER")
 263    homologs: (
 264        list
 265        | Annotated[
 266            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 267            Field(discriminator="kind"),
 268        ]
 269    ) = Field(description="The homologous assemblies used for training.")
 270    dataset: (
 271        list
 272        | Annotated[
 273            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 274            Field(discriminator="kind"),
 275        ]
 276        | None
 277    ) = Field(
 278        description="The assayed assemblies used for training. If None, predictor finetuning and generator conditioning will be skipped."
 279    )
 280
 281
 282class MonomerHomologsData(BaseModel):
 283    """Monomer homolog sequences for base model training."""
 284
 285    kind: Literal["MONOMER"] = Field(default="MONOMER")
 286    homologs: (
 287        list
 288        | Annotated[
 289            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 290            Field(discriminator="kind"),
 291        ]
 292    ) = Field(
 293        description="Homologous sequences used for training. These sequences should be evolutionarily related to the target protein family."
 294    )
 295
 296
 297class MonomerModels(BaseModel):
 298    kind: Literal["MONOMER"] = Field(default="MONOMER")
 299    sampler: ArtifactParam = Field(description="The monomer sampler model.")
 300    scorer: ArtifactParam | None = Field(description="The monomer scorer model.")
 301    conditioned_sampler: ArtifactParam | None = Field(
 302        description="The monomer sampler model conditioned on the primary objective."
 303    )
 304
 305
 306class MonomerSamplerConditioningData(BaseModel):
 307    kind: Literal["MONOMER"] = Field(default="MONOMER")
 308    train_dataset: (
 309        list
 310        | Annotated[
 311            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 312            Field(discriminator="kind"),
 313        ]
 314    ) = Field(description="The training dataset containing assayed assemblies.")
 315    val_dataset: (
 316        list
 317        | Annotated[
 318            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 319            Field(discriminator="kind"),
 320        ]
 321        | None
 322    ) = Field(
 323        default=None,
 324        description="The validation dataset containing assayed assemblies. If not set, the training dataset will be split to create a validation dataset.",
 325    )
 326    sampler: ArtifactParam = Field(description="The pre-trained sampler model to condition.")
 327
 328
 329class MonomerSamplerConditioningResult(BaseModel):
 330    kind: Literal["MONOMER"] = Field(default="MONOMER")
 331    conditioned_sampler: ArtifactParam = Field(description="The conditioned sampler model.")
 332
 333
 334class MonomerSamplerConfig(BaseModel):
 335    """Configuration for a monomer sampler."""
 336
 337    kind: Literal["MONOMER_SAMPLER"] = Field(default="MONOMER_SAMPLER")
 338    sampler: ArtifactParam = Field(description="The trained sampler artifact used to compute sequence likelihoods.")
 339    reference_assembly: MonomerAssembly = Field(
 340        description="The reference assembly used as the denominator when computing sequence likelihood ratios."
 341    )
 342    controls: (
 343        list
 344        | Annotated[
 345            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 346            Field(discriminator="kind"),
 347        ]
 348        | None
 349    ) = Field(
 350        default=None,
 351        description="Optional control assemblies to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_assemblies` will be intentionally diverse from these assemblies. However these assemblies will *not* be included in the `selected_assemblies` output table, so -- if used -- take care to re-include them. ",
 352    )
 353
 354
 355class MonomerSamplerInput(BaseModel):
 356    kind: Literal["MONOMER"] = Field(default="MONOMER")
 357    assemblies: (
 358        list
 359        | Annotated[
 360            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 361            Field(discriminator="kind"),
 362        ]
 363    ) = Field(description="Monomer assemblies to run inference on.")
 364    sampler: ArtifactParam = Field(description="The trained sampler model.")
 365    reference_assembly: MonomerAssembly | None = Field(
 366        default=None,
 367        description="Optional reference assembly for computing pseudo log-likelihood ratios. All assemblies must be the same length as the reference assembly. When omitted, raw pseudo log-likelihoods are returned.",
 368    )
 369
 370
 371class MonomerSamplingTemplate(BaseModel):
 372    kind: Literal["MONOMER"] = Field(default="MONOMER")
 373    assembly: MonomerAssembly = Field(description="The template sequence to use for generation.")
 374    sampler: ArtifactParam = Field(description="The trained sampler artifact used to generate candidate sequences.")
 375
 376
 377class MonomerScorerConfig(BaseModel):
 378    """Configuration for a monomer scorer."""
 379
 380    kind: Literal["MONOMER_SCORER"] = Field(default="MONOMER_SCORER")
 381    scorer: ArtifactParam = Field(description="The trained scoring model artifact used to score sequences.")
 382    controls: (
 383        list
 384        | Annotated[
 385            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 386            Field(discriminator="kind"),
 387        ]
 388        | None
 389    ) = Field(
 390        default=None,
 391        description="Optional control sequences to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_sequences` will be intentionally diverse from these sequences. However these sequences will *not* be included in the `selected_sequences` output table, so -- if used -- take care to re-include them. ",
 392    )
 393
 394
 395class MonomerScorerInput(BaseModel):
 396    kind: Literal["MONOMER"] = Field(default="MONOMER")
 397    assemblies: (
 398        list
 399        | Annotated[
 400            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 401            Field(discriminator="kind"),
 402        ]
 403    ) = Field(description="Monomer assemblies to run inference on.")
 404    scorer: ArtifactParam | CustomPredictorParam = Field(description="The trained scorer ensemble.")
 405    noise: SamplingNoise | None = Field(
 406        default=None,
 407        description="MC-dropout noise configuration. When omitted, returns un-noised deterministic scores.",
 408    )
 409
 410
 411class MonomerScorerModel(BaseModel):
 412    kind: Literal["MONOMER"] = Field(default="MONOMER")
 413    scorer: ArtifactParam = Field(description="The trained scorer model for monomer sequences.")
 414
 415
 416class MonomerScorerRef(BaseModel):
 417    """A reference to a monomer scorer model."""
 418
 419    kind: Literal["MONOMER"] = Field(default="MONOMER")
 420    scorer: ArtifactParam | CustomPredictorParam = Field(
 421        description="Monomer scorer model. Can be a reference to a custom predictor."
 422    )
 423
 424
 425class MonomerSource(BaseModel):
 426    """A source of monomer assemblies."""
 427
 428    kind: Literal["MONOMER"] = Field(default="MONOMER")
 429    source: (
 430        list
 431        | Annotated[
 432            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 433            Field(discriminator="kind"),
 434        ]
 435    )
 436
 437
 438class MonomerTrainBaseModelResult(BaseModel):
 439    """Result containing a trained monomer sampler."""
 440
 441    kind: Literal["MONOMER"] = Field(default="MONOMER")
 442    sampler: ArtifactParam = Field(description="The trained sampler model that can be used for sequence generation.")
 443
 444
 445class MonomerTrainScorerInputs(BaseModel):
 446    kind: Literal["MONOMER"] = Field(default="MONOMER")
 447    dataset: (
 448        list
 449        | Annotated[
 450            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 451            Field(discriminator="kind"),
 452        ]
 453    ) = Field(description="The dataset containing assayed assemblies for training the scorer.")
 454    base_model: ArtifactParam = Field(description="The pre-trained base model (sampler) to use for scorer training.")
 455
 456
 457class OASDatabaseSearch(BaseModel):
 458    type_: Literal["OAS_SEARCH"] = Field(default="OAS_SEARCH", alias="type")
 459    seed_assemblies: list[VhVlAssembly] = Field(
 460        description="The list of light/heavy chain assemblies to search for.", min_length=1, max_length=100
 461    )
 462    search_unpaired: bool = Field(default=True, description="Whether to search for unpaired heavy/light chains.")
 463    search_paired: bool = Field(default=True, description="Whether to search for paired heavy/light chains.")
 464    annotate_assembly_features: bool = Field(
 465        default=True, description="Whether to infer CDR features based on OAS annotations."
 466    )
 467
 468
 469class ObjectivesAndConstraints(BaseModel):
 470    """Evaluation via a primary objective and/or constraints.
 471
 472    At least one of ``primary_objective`` or ``constraints`` must be provided. The output
 473    property ordering is ``[primary_objective, *constraints]`` when a primary objective is
 474    present; otherwise ``[*constraints]``.
 475    """
 476
 477    primary_objective: PrimaryObjective | None = Field(
 478        default=None,
 479        description="The primary objective defining the optimisation direction and reference value. Used to compute reference-subtracted scores and normalise direction (higher = better).",
 480    )
 481    constraints: list[Constraint] = Field(
 482        default_factory=list, description="Constraints to include in the score computation.", max_length=50
 483    )
 484
 485
 486class Parallel(BaseModel):
 487    """An ensemble of generators to use in parallel.
 488
 489    All generators in the ensemble run independently and each contribute an equal
 490    number of candidate assemblies for evaluation. This does not guarantee that each
 491    generator will contribute an equal number of assemblies to `selected_assemblies`,
 492    as assemblies from some generators in `ensemble` may perform better than others.
 493    """
 494
 495    type: Literal["PARALLEL"] = Field(default="PARALLEL")
 496    ensemble: list[
 497        Annotated[
 498            Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator,
 499            Field(description="The generation strategy for creating candidate sequences.", discriminator="type"),
 500        ]
 501    ] = Field(default_factory=list, description="The list of generators to run in parallel within this ensemble.")
 502    alias: str | None = Field(
 503        default=None,
 504        description="Optional alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. If provided, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
 505    )
 506
 507
 508class PreRanked(BaseModel):
 509    """Configuration for selecting from pre-ranked assemblies.
 510
 511    With this option, the assemblies passed for selection are assumed to be ranked in descending
 512    order of desirability. Selection is based on this order and subject to diversity constraints that
 513    can be controlled via the `hit_redundancy` parameter.
 514    """
 515
 516    type: Literal["PRE_RANKED"] = Field(default="PRE_RANKED")
 517    hit_redundancy: int = Field(
 518        default=3,
 519        description='The maximum number of similar variants within the proposed set.\n\nWe say two variants are "similar" when their predicted performances are highly correlated.\nAt lower values this will mean on average fewer hits, but greater diversity within the set,\nwhile at maximum value (`hit_redundancy` equals the number of assemblies to select) the selection is purely greedy, with\nno requirement on diversity within the set.',
 520    )
 521
 522
 523class PrimaryObjective(BaseModel):
 524    """The primary objective of an optimization task."""
 525
 526    measure: Measure = Field(description="The measure to optimize.")
 527    reference: float | RelativeTo = Field(
 528        description="The reference relative to which the primary objective is defined. For multiplicative assays, an absolute reference has to be positive."
 529    )
 530
 531
 532class RangeAnnotation(BaseModel):
 533    ranges: dict[str, list[tuple[int, int]]] = Field(
 534        description="A mapping from polymer names to annotation ranges. Each range is a tuple of (start, end). Indices are 0-based and inclusive for the start and exclusive for the end. For example, (0, 5) refers to positions 0 to 4."
 535    )
 536
 537
 538class RelativeTo(BaseModel):
 539    """A constraint or objective relative to a reference assembly."""
 540
 541    reference: MonomerAssembly | VhVlAssembly = Field(description="The reference assembly.")
 542    margin: float | None = Field(
 543        default=None,
 544        description="The desired margin relative to the `reference` assembly.\n\n        Denoting the score of the `reference` assembly as `reference_score`, the effective threshold will be `reference_score + margin` for additive assays\n        and `reference_score * margin` for multiplicative assays. Note: for multiplicative assays, the `margin` must be positive.\n\n        As an example, to tighten a constraint on a multiplicative assay with direction `MAXIMIZE`, specify a `margin`\n        greater than 1.0. Alternatively, if the direction is `MINIMIZE`, specify a `margin` less than 1.0.\n        ",
 545    )
 546
 547
 548class Requirement(BaseModel):
 549    """Assembly generation requirements specifying required and forbidden features on the generated assemblies.
 550
 551    Each requirement defines constraints for specific motifs in specific ranges of the assembly.
 552    The `operator` field determines how motif-range pairs are combined:
 553    - operator="all": *All* motif-range pairs must match
 554    - operator="any": *At least one* motif-range pair must match
 555    - operator="none": *No* motif-range pairs must match
 556
 557    When multiple requirements are provided in a list, *all* requirements must be satisfied.
 558
 559    Examples:
 560    1. Block mutations to Methionine ("M") across the entire heavy and light chains to prevent oxidation:
 561
 562       Requirement(
 563           motifs=["M"],
 564           apply_to=RangeAnnotation(ranges={"vh": [(0, 120)], "vl": [(0, 110)]}),
 565           operator="none",
 566       )
 567
 568    2. Avoid deamidation motifs ("NG", "NS", "NT") in the heavy chain:
 569
 570       Requirement(
 571           motifs=["NG", "NS", "NT"],
 572           apply_to=RangeAnnotation(ranges={"vh": [(0, 120)]}),
 573           operator="none",
 574       )
 575
 576    3. Ensure Cysteines are present at positions 22 and 96 in the heavy chain:
 577
 578       Requirement(
 579           motifs=["C"],
 580           apply_to=RangeAnnotation(ranges={"vh": [(22, 23), (96, 97)]}),
 581           operator="all",
 582       )
 583    """
 584
 585    apply_to: RangeAnnotation
 586    motifs: list[Annotated[str, Field(description="A motif or regular expression of amino acids.")]]
 587    operator: Literal["none", "all", "any"] = Field(
 588        description="The choice of how to interpret the motif-range pairs as a single requirement."
 589    )
 590
 591
 592class SamplerConditioningParameters(ParametersBase):
 593    task_type: ClassVar[str] = "condition_sampler/v2"
 594
 595    data: MonomerSamplerConditioningData | VhVlSamplerConditioningData = Field(
 596        description="The training and validation data along with the pre-trained sampler model.", discriminator="kind"
 597    )
 598    primary_objective: Measure = Field(
 599        description="The assay on which to condition the sampler model. Must correspond to an assay present in the datasets."
 600    )
 601
 602
 603class SamplerConditioningResult(ResultBase):
 604    models: MonomerSamplerConditioningResult | VhVlSamplerConditioningResult = Field(
 605        description="The conditioned sampler.", discriminator="kind"
 606    )
 607    report: ArtifactParam | None = Field(default=None, description="A report on the result of the task.")
 608
 609
 610class SamplerInferenceParameters(ParametersBase):
 611    """Parameters for the sampler inference endpoint.
 612
 613    Runs a trained sampler over a set of sequences and emits per-sequence sampler
 614    scores into the ``cradle.inferences.sampler_predictions`` BigQuery table.
 615    """
 616
 617    task_type: ClassVar[str] = "inference_sampler/v2"
 618
 619    sampler_input: MonomerSamplerInput | VhVlSamplerInput = Field(
 620        description="Assemblies and sampler model to use for inference.", discriminator="kind"
 621    )
 622    hyperparameters: dict[str, Any] = Field(
 623        default_factory=dict,
 624        description="A dictionary of hyperparameter overrides to apply to the default hyperparameters used by the task.",
 625    )
 626
 627
 628class SamplerInferenceResult(ResultBase):
 629    pass
 630
 631
 632class SamplingNoise(BaseModel):
 633    """Configuration for MC-dropout score sampling. When omitted on the scorer input, un-noised scores are returned."""
 634
 635    num_samples: int = Field(
 636        default=64, description="Number of score samples to draw per ensemble member and property."
 637    )
 638    dropout_p: float = Field(default=0.3, description="Dropout probability used to generate score samples.")
 639    noise_seed: int | None = Field(
 640        default=None,
 641        description="Seed for noise generation. If None, a fixed default seed is used for reproducibility.",
 642    )
 643
 644
 645class ScoreBasedFilterer(BaseModel):
 646    """Filters assemblies based on whether predicted performance meets constraints.
 647
 648    The scorer evaluates assemblies and only those that meet all specified constraints are retained.
 649    This can, for example, be used to incorporate in-silico metrics in assembly generation.
 650    """
 651
 652    kind: Literal["SCORED_FILTERER"] = Field(default="SCORED_FILTERER")
 653    scorer: MonomerScorerRef | VhVlScorerRef = Field(
 654        description="The scorer used to evaluate assemblies for filtering.", discriminator="kind"
 655    )
 656    constraints: list[Constraint] = Field(
 657        default_factory=list, description="List of constraints to apply during filtering.", max_length=50
 658    )
 659
 660
 661class ScoreBasedRanker(BaseModel):
 662    """Configuration for ranking and selecting generated assemblies."""
 663
 664    type: Literal["SCORE_BASED"] = Field(default="SCORE_BASED")
 665    primary_objective: PrimaryObjective = Field(description="The primary optimization objective of `engineer`.")
 666    constraints: list[Constraint] = Field(
 667        default_factory=list, description="List of constraints `engineer` will attempt to satisfy.", max_length=50
 668    )
 669    scorer_config: MonomerScorerConfig | VhVlScorerConfig = Field(
 670        description="The trained scoring model artifact used to score sequences, along with any optional control sequences to include in the final output for experimental validation.",
 671        discriminator="kind",
 672    )
 673    hit_redundancy: int | None = Field(
 674        default=None,
 675        description='The maximum number of similar variants within the proposed set.\n\nWe say two variants are "similar" when their predicted performances are highly correlated.\nAt lower values this will mean on average fewer hits, but greater diversity within the set,\nwhile at maximum value (`hit_redundancy` equals the number of assemblies to select) the selection is purely greedy, with\nno requirement on diversity within the set. If set to None, it will be automatically set\nto a default value of `1 + num_assemblies//20`.',
 676    )
 677
 678
 679class ScorerInferenceParameters(ParametersBase):
 680    """Parameters for the scorer inference endpoint.
 681
 682    Runs inference using a trained scorer ensemble and emits per-sequence scores
 683    into ``cradle.inferences.scorer_predictions``. When a noise configuration is provided
 684    on the scorer input, produces MC-dropout score samples; otherwise returns un-noised
 685    deterministic scores (``num_samples=1``). For each property the output is a flat array
 686    of shape ``[num_models, num_samples]`` representing either the predicted value above the
 687    property's reference, in direction-normalised log space (if `score_by` is provided) or
 688    the raw scorer values otherwise.
 689    """
 690
 691    task_type: ClassVar[str] = "inference_scorer/v2"
 692
 693    scorer_input: MonomerScorerInput | VhVlScorerInput = Field(
 694        description="Assemblies and scorer model to use for inference.", discriminator="kind"
 695    )
 696    score_by: ObjectivesAndConstraints | None = Field(
 697        description="Defines the properties to evaluate: a primary objective and/or constraints."
 698    )
 699    hyperparameters: dict[str, Any] = Field(
 700        default_factory=dict,
 701        description="A dictionary of hyperparameter overrides to apply to the default hyperparameters used by the task.",
 702    )
 703
 704
 705class ScorerInferenceResult(ResultBase):
 706    pass
 707
 708
 709class SearchParameters(ParametersBase):
 710    task_type: ClassVar[str] = "search/v2"
 711
 712    source: UnirefDatabaseSearch | OASDatabaseSearch = Field(
 713        description="The source database to search against, along with the assemblies to search for, and database-specific configuration.",
 714        discriminator="type_",
 715    )
 716
 717
 718class SearchResult(ResultBase):
 719    assembly_features: list[dict[str, AssemblyFeature]] = Field(
 720        default_factory=list,
 721        description="Parallel array (wrt to `seed_assemblies`) of features, keyed by feature name.",
 722    )
 723    homologs: TableResult
 724    report: ArtifactParam | None = Field(default=None, description="A report on the result of the task.")
 725
 726
 727class SelectParameters(ParametersBase):
 728    """Parameters for the select task."""
 729
 730    task_type: ClassVar[str] = "select/v2"
 731
 732    assemblies: MonomerSource | VhVlSource = Field(description="The source of assemblies to select from.")
 733    ranker: (
 734        Annotated[
 735            LikelihoodBasedRanker | ScoreBasedRanker,
 736            Field(description="The ranking strategy for selecting candidate assemblies.", discriminator="type"),
 737        ]
 738        | PreRanked
 739    ) = Field(
 740        description="The configuration defining how candidate assemblies are evaluated for inclusion in `selected_assemblies`."
 741    )
 742    min_assemblies_to_select: int = Field(
 743        description="The minimum number of assemblies to select. The number of selected candidates is dynamically determined and may be greater than this number. If an exact number of assemblies is desired, set `min_assemblies_to_select` and `max_assemblies_to_select` to the same value."
 744    )
 745    max_assemblies_to_select: int | None = Field(
 746        default=None,
 747        description="An optional (inclusive) upper limit on the number of assemblies to select. If provided, the task will select between `min_assemblies_to_select` and `max_assemblies_to_select` assemblies.",
 748    )
 749
 750
 751class SelectResult(ResultBase):
 752    """Result of the select task."""
 753
 754    selected_assemblies: TableResult = Field(
 755        description="Table containing the selected assemblies output by `select` sorted in descending order of preference."
 756    )
 757
 758
 759class Sequential(BaseModel):
 760    """An ensemble of generators to evaluate sequentially.
 761
 762    Each generator will have the opportunity to evaluate an equal number of candidate
 763    assemblies. However, if earlier generators are able to satisfy the requirements of
 764    the `ranker`, later generators may never be used. In this way this method of
 765    ensembling expresses a strict preference among the generators in `ensemble`.
 766    """
 767
 768    type: Literal["SEQUENTIAL"] = Field(default="SEQUENTIAL")
 769    ensemble: list[
 770        Annotated[
 771            Sequential | Parallel | ModelBasedGenerator | Explicit | CombinatorialGenerator,
 772            Field(description="The generation strategy for creating candidate sequences.", discriminator="type"),
 773        ]
 774    ] = Field(default_factory=list, description="The list of generators to run in sequence within this ensemble.")
 775    alias: str | None = Field(
 776        default=None,
 777        description="Optional alias to refer to this generator within the task. Useful for tracking provenance of assemblies generated by this generator. If provided, this name will appear in the `generator_provenance` column for rows in `selected_sequences` which were produced by this generator.",
 778    )
 779
 780
 781class TrainBaseModelParameters(ParametersBase):
 782    """Parameters for the base model training (evotuning) task.
 783
 784    This task fine-tunes a pretrained language model on a set of homologous sequences using masked language modeling.
 785    The resulting model can be used as a sampler for downstream generation tasks.
 786    """
 787
 788    task_type: ClassVar[str] = "train_base_model/v2"
 789
 790    data: MonomerHomologsData | VhVlHomologsData = Field(
 791        description="Homologous sequences for training the base model. The model learns the sequence distribution of the provided protein family.",
 792        discriminator="kind",
 793    )
 794
 795
 796class TrainBaseModelResult(ResultBase):
 797    """Result of the base model training task."""
 798
 799    models: MonomerTrainBaseModelResult | VhVlTrainBaseModelResult = Field(
 800        description="The trained sampler model.", discriminator="kind"
 801    )
 802    report: ArtifactParam | None = Field(default=None, description="Training report with metrics and diagnostics.")
 803
 804
 805class TrainParameters(ParametersBase):
 806    task_type: ClassVar[str] = "train/v2"
 807
 808    data: MonomerData | VhVlData = Field(
 809        description="The training data, including homologs and assayed sequences.", discriminator="kind"
 810    )
 811    assays: list[Assay] = Field(
 812        default_factory=list,
 813        description="The assays used for sampler conditioning (on the `primary_objective`) and scorer finetuning and evaluation. Must correspond to assays present in `data.dataset`. If `data.dataset` is None, `assays` must be empty. Otherwise, it must contain at least one assay.",
 814        max_length=50,
 815    )
 816    primary_objective: Measure | None = Field(
 817        default=None,
 818        description="The assay on which the sampler model is conditioned. Should correspond to an assay in `dataset`. If `data.dataset` is None, `primary_objective` must not be set. Otherwise, it must be set.",
 819    )
 820
 821
 822class TrainResult(ResultBase):
 823    models: MonomerModels | VhVlModels = Field(description="The trained models.", discriminator="kind")
 824    report: ArtifactParam | None = Field(description="A report on the result of the task.")
 825
 826
 827class TrainScorerParameters(ParametersBase):
 828    task_type: ClassVar[str] = "train_scorer/v2"
 829
 830    data: MonomerTrainScorerInputs | VhVlTrainScorerInputs = Field(
 831        description="The training data and pre-trained base model for scorer training.", discriminator="kind"
 832    )
 833    assays: list[Assay] = Field(
 834        default_factory=list,
 835        description="The assays used for scorer training and evaluation. Must correspond to assays present in `data.dataset`. Must contain at least one assay.",
 836        max_length=50,
 837    )
 838
 839
 840class TrainScorerResult(ResultBase):
 841    scorer_model: MonomerScorerModel | VhVlScorerModel = Field(
 842        description="The trained scorer model.", discriminator="kind"
 843    )
 844    report: ArtifactParam | None = Field(default=None, description="A report on the result of the task.")
 845
 846
 847class UnirefDatabaseSearch(BaseModel):
 848    type_: Literal["UNIREF_SEARCH"] = Field(default="UNIREF_SEARCH", alias="type")
 849    seed_assemblies: list[MonomerAssembly] = Field(
 850        description="A list of monomeric assemblies to search for.", min_length=1, max_length=100
 851    )
 852
 853
 854class VhVlAssembly(BaseModel):
 855    """An assembly of a VH and a VL chain."""
 856
 857    vh: str | None = Field(default=None, description="The heavy chain sequence.")
 858    vl: str | None = Field(default=None, description="The light chain sequence.")
 859
 860
 861class VhVlData(BaseModel):
 862    kind: Literal["VHVL"] = Field(default="VHVL")
 863    homologs: (
 864        list
 865        | Annotated[
 866            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 867            Field(discriminator="kind"),
 868        ]
 869    ) = Field(description="The homologous assemblies used for training.")
 870    dataset: (
 871        list
 872        | Annotated[
 873            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 874            Field(discriminator="kind"),
 875        ]
 876        | None
 877    ) = Field(
 878        description="The assayed assemblies used for training. If None, predictor finetuning and generator conditioning will be skipped."
 879    )
 880
 881
 882class VhVlHomologsData(BaseModel):
 883    """VH/VL homolog sequences for base model training."""
 884
 885    kind: Literal["VHVL"] = Field(default="VHVL")
 886    homologs: (
 887        list
 888        | Annotated[
 889            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 890            Field(discriminator="kind"),
 891        ]
 892    ) = Field(
 893        description="Homologous antibody sequences used for training. These sequences should be evolutionarily related to the target antibody family."
 894    )
 895
 896
 897class VhVlModels(BaseModel):
 898    kind: Literal["VH_VL"] = Field(default="VH_VL")
 899    sampler: ArtifactParam = Field(description="The VhVl sampler model.")
 900    scorer: ArtifactParam | None = Field(description="The VhVl scorer model.")
 901    conditioned_sampler: ArtifactParam | None = Field(
 902        description="The VhVl sampler model conditioned on the primary objective."
 903    )
 904
 905
 906class VhVlSamplerConditioningData(BaseModel):
 907    kind: Literal["VH_VL"] = Field(default="VH_VL")
 908    train_dataset: (
 909        list
 910        | Annotated[
 911            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 912            Field(discriminator="kind"),
 913        ]
 914    ) = Field(description="The training dataset containing assayed assemblies.")
 915    val_dataset: (
 916        list
 917        | Annotated[
 918            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 919            Field(discriminator="kind"),
 920        ]
 921        | None
 922    ) = Field(
 923        default=None,
 924        description="The validation dataset containing assayed assemblies. If not set, the training dataset will be split to create a validation dataset.",
 925    )
 926    sampler: ArtifactParam = Field(description="The pre-trained sampler model to condition.")
 927
 928
 929class VhVlSamplerConditioningResult(BaseModel):
 930    kind: Literal["VH_VL"] = Field(default="VH_VL")
 931    conditioned_sampler: ArtifactParam = Field(description="The conditioned sampler model.")
 932
 933
 934class VhVlSamplerConfig(BaseModel):
 935    """Configuration for a VhVl sampler."""
 936
 937    kind: Literal["VHVL_SAMPLER"] = Field(default="VHVL_SAMPLER")
 938    sampler: ArtifactParam = Field(description="The trained sampler artifact used to compute assembly likelihoods.")
 939    reference_assembly: VhVlAssembly = Field(
 940        description="The reference assembly used as the denominator when computing assembly likelihood ratios."
 941    )
 942    controls: (
 943        list
 944        | Annotated[
 945            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 946            Field(discriminator="kind"),
 947        ]
 948        | None
 949    ) = Field(
 950        default=None,
 951        description="Optional control assemblies to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_assemblies` will be intentionally diverse from these assemblies. However these assemblies will *not* be included in the `selected_assemblies` output table, so -- if used -- take care to re-include them. ",
 952    )
 953
 954
 955class VhVlSamplerInput(BaseModel):
 956    kind: Literal["VHVL"] = Field(default="VHVL")
 957    assemblies: (
 958        list
 959        | Annotated[
 960            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 961            Field(discriminator="kind"),
 962        ]
 963    ) = Field(description="VH/VL assemblies to run inference on.")
 964    sampler: ArtifactParam = Field(description="The trained sampler model.")
 965    reference_assembly: VhVlAssembly | None = Field(
 966        default=None,
 967        description="Optional reference assembly for computing pseudo log-likelihood ratios. All assemblies must be the same length as the reference assembly. When omitted, raw pseudo log-likelihoods are returned.",
 968    )
 969
 970
 971class VhVlSamplingTemplate(BaseModel):
 972    kind: Literal["VHVL"] = Field(default="VHVL")
 973    assembly: VhVlAssembly = Field(description="The template sequence to use for generation.")
 974    sampler: ArtifactParam = Field(description="The trained sampler artifact used to generate candidate sequences.")
 975
 976
 977class VhVlScorerConfig(BaseModel):
 978    """Configuration for a VhVl scorer."""
 979
 980    kind: Literal["VHVL_SCORER"] = Field(default="VHVL_SCORER")
 981    scorer: ArtifactParam = Field(description="The trained scoring model artifact used to score sequences.")
 982    controls: (
 983        list
 984        | Annotated[
 985            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 986            Field(discriminator="kind"),
 987        ]
 988        | None
 989    ) = Field(
 990        default=None,
 991        description="Optional control sequences to include in the final output for experimental validation. These controls should generally be included in the assayed plate, as `selected_sequences` will be intentionally diverse from these sequences. However these sequences will *not* be included in the `selected_sequences` output table, so -- if used -- take care to re-include them. ",
 992    )
 993
 994
 995class VhVlScorerInput(BaseModel):
 996    kind: Literal["VHVL"] = Field(default="VHVL")
 997    assemblies: (
 998        list
 999        | Annotated[
1000            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
1001            Field(discriminator="kind"),
1002        ]
1003    ) = Field(description="VH/VL assemblies to run inference on.")
1004    scorer: ArtifactParam | CustomPredictorParam = Field(
1005        description="VH/VL  scorer model. Can be a referenceto a custom predictor."
1006    )
1007    noise: SamplingNoise | None = Field(
1008        default=None,
1009        description="MC-dropout noise configuration. When omitted, returns un-noised deterministic scores.",
1010    )
1011
1012
1013class VhVlScorerModel(BaseModel):
1014    kind: Literal["VH_VL"] = Field(default="VH_VL")
1015    scorer: ArtifactParam = Field(description="The trained scorer model for antibody sequences.")
1016
1017
1018class VhVlScorerRef(BaseModel):
1019    """A reference to a VH/VL scorer model."""
1020
1021    kind: Literal["VHVL"] = Field(default="VHVL")
1022    scorer: ArtifactParam | CustomPredictorParam = Field(
1023        description="VH/VL scorer model. Can be a reference to a custom predictor."
1024    )
1025
1026
1027class VhVlSource(BaseModel):
1028    """A source of VH/VL assemblies."""
1029
1030    kind: Literal["VHVL"] = Field(default="VHVL")
1031    source: (
1032        list
1033        | Annotated[
1034            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
1035            Field(discriminator="kind"),
1036        ]
1037    )
1038
1039
1040class VhVlTrainBaseModelResult(BaseModel):
1041    """Result containing a trained VH/VL sampler."""
1042
1043    kind: Literal["VHVL"] = Field(default="VHVL")
1044    sampler: ArtifactParam = Field(
1045        description="The trained sampler model that can be used for antibody sequence generation."
1046    )
1047
1048
1049class VhVlTrainScorerInputs(BaseModel):
1050    kind: Literal["VHVL"] = Field(default="VHVL")
1051    dataset: (
1052        list
1053        | Annotated[
1054            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
1055            Field(discriminator="kind"),
1056        ]
1057    ) = Field(description="The dataset containing assayed assemblies for training the scorer.")
1058    base_model: ArtifactParam = Field(description="The pre-trained base model (sampler) to use for scorer training.")