cradle.sdk.types.tasks.v1

  1from __future__ import annotations
  2
  3from enum import StrEnum
  4from typing import Annotated, 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 FeatureAnnotation(StrEnum):
 23    """Annotation to be configured per residue in a domain:
 24
 25    UNRELIABLE: The evolutionary context is not reliable, eg in CDRs of antibodies.
 26    IGNORE: A part of the sequence that should not be used for evolutionary search and won't be mutated, this can be tags, linkers or other artificial subsequences.
 27    """
 28
 29    UNRELIABLE = "UNRELIABLE"
 30    IGNORE = "IGNORE"
 31
 32
 33class OptimizationDirection(StrEnum):
 34    MAXIMIZE = "MAXIMIZE"
 35    MINIMIZE = "MINIMIZE"
 36
 37
 38class ProteinDatabase(StrEnum):
 39    UNIREF_30 = "UNIREF_30"
 40    OAS_90 = "OAS_90"
 41
 42
 43class ProteinModelType(StrEnum):
 44    DEFAULT = "DEFAULT"
 45    ANTIBODY = "ANTIBODY"
 46
 47
 48class ScaleType(StrEnum):
 49    MULTIPLICATIVE = "MULTIPLICATIVE"
 50    ADDITIVE = "ADDITIVE"
 51    RANK = "RANK"
 52
 53
 54class AbsoluteConstraint(BaseModel):
 55    type_: Literal["ABSOLUTE_CONSTRAINT"] = Field(default="ABSOLUTE_CONSTRAINT", alias="type")
 56    assay_id: str = Field(description="The ID of the constrained assay.")
 57    direction: OptimizationDirection = Field(
 58        description="Whether the assay value should be above (MAXIMIZE) or below (MINIMIZE) the threshold."
 59    )
 60    threshold: float = Field(description="The threshold assay value.")
 61
 62
 63class AnalyzeDataParameters(ParametersBase):
 64    task_type: ClassVar[str] = "analyze.data/v1"
 65
 66    reference_sequence: str = Field(description="The sequence to consider as a reference for computing mutations.")
 67    dataset: (
 68        list
 69        | Annotated[
 70            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 71            Field(discriminator="kind"),
 72        ]
 73    ) = Field(description="The assay data to be used for model training.")
 74    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries", max_length=50)
 75
 76
 77class AnalyzeDiversifyParameters(ParametersBase):
 78    task_type: ClassVar[str] = "analyze.diversify/v1"
 79
 80    selected_sequences: (
 81        list
 82        | Annotated[
 83            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 84            Field(discriminator="kind"),
 85        ]
 86    ) = Field(description="The name of the input table with the selected sequences.")
 87    generated_sequences: (
 88        list
 89        | Annotated[
 90            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
 91            Field(discriminator="kind"),
 92        ]
 93    ) = Field(description="The name of the input view/table with the generated sequences.")
 94    template_metadata: TemplateMetadata = Field(description="Metadata about templates.")
 95
 96
 97class AnalyzeEngineerParameters(ParametersBase):
 98    task_type: ClassVar[str] = "analyze.engineer/v1"
 99
100    selected_sequences: (
101        list
102        | Annotated[
103            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
104            Field(discriminator="kind"),
105        ]
106    ) = Field(description="The name of the input table with the selected sequences.")
107    engineered_sequences: (
108        list
109        | Annotated[
110            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
111            Field(discriminator="kind"),
112        ]
113    ) = Field(description="The name of the input view/table with the generated sequences.")
114    assays: list[Assay] = Field(
115        default_factory=list, description="The list of objectives that were optimized during training.", max_length=50
116    )
117    primary_objective: PrimaryObjective = Field(description="The primary objective to be optimized.")
118    constraints: list[Annotated[AbsoluteConstraint | RelativeConstraint, Field(discriminator="type_")]] = Field(
119        default_factory=list, description="The constraints applied to engineering.", max_length=50
120    )
121    template_metadata: list[TemplateMetadata] = Field(description="Metadata about templates.")
122    new_mutation_ratio: float = Field(description="Ratio of new mutations in the generated sequences.")
123    scorer: ArtifactParam = Field(description="Scoring model used to predict assay values for candidate sequences.")
124    custom_predictor_assays: list[Assay] = Field(
125        default_factory=list, description="The list of assays predicted by custom predictors.", max_length=50
126    )
127
128
129class AnalyzeResult(ResultBase):
130    report: ArtifactParam = Field(description="The report data")
131
132
133class AnalyzeTrainParameters(ParametersBase):
134    task_type: ClassVar[str] = "analyze.train/v1"
135
136    prediction_data: (
137        list
138        | Annotated[
139            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
140            Field(discriminator="kind"),
141        ]
142    ) = Field(description="Table of predicted assay values for the input dataset under the scorer.")
143    generator_prediction_data: (
144        list
145        | Annotated[
146            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
147            Field(discriminator="kind"),
148        ]
149    ) = Field(description="Pseudo log-likelihoods of the sequences in the dataset under the sampler.")
150    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries", max_length=50)
151    primary_objective: PrimaryObjective = Field(
152        description="The primary objective on which the sampler model is conditioned."
153    )
154    constraints: list[Annotated[AbsoluteConstraint | RelativeConstraint, Field(discriminator="type_")]] = Field(
155        default_factory=list, description="The constraints applied to training.", max_length=50
156    )
157    custom_predictors: list[CustomPredictorParam] = Field(
158        default_factory=list,
159        description="Custom predictors used to score input sequences, enabling comparison with trained models.",
160        max_length=10,
161    )
162
163
164class Assay(BaseModel):
165    """Metadata of an assay."""
166
167    assay_id: str = Field(description="ID of the assay", min_length=1, max_length=124)
168    name: str = Field(description="Human readable name of the assay", min_length=1, max_length=100)
169    scale_type: ScaleType = Field(
170        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."
171    )
172    unit: str | None = Field(
173        default=None, description="Unit of the assay - used for display purposes only", min_length=1, max_length=100
174    )
175
176
177class BlockedAAItem(BaseModel):
178    blocked_aas: str = Field(
179        description="One or more amino acid letter code to be blocked from being mutated to. Use '*' to indicate that a position is fully blocked from being mutated."
180    )
181    ranges: list[tuple[int, int]] = Field(
182        description="List of ranges where the mutations are blocked. 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) blocks positions 0 to 4."
183    )
184
185
186class BlockedMotifItem(BaseModel):
187    blocked_motif: str = Field(description="Amino acid motif to be blocked from appearing in the final sequence.")
188    ranges: list[tuple[int, int]] = Field(
189        description="List of ranges where the motif is blocked. 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) blocks positions 0 to 4. The entire motif must fit into the specified positions to be blocked."
190    )
191
192
193class ClusterParameters(BaseModel):
194    """Parameters for clustering of sequences."""
195
196    e_value_threshold: float = Field(
197        default=float("inf"),
198        description="Maximum e-value allowed between linked sequences in clustering algorithm. A lower number will result in tighter clusters.",
199    )
200    coverage_threshold: float = Field(
201        default=0.8,
202        description="Minimum coverage allowed between linked sequences in clustering algorithm. A higher number will result in tighter clusters.",
203    )
204    sequence_identity_threshold: float = Field(
205        default=0.0,
206        description="Minimum sequence identity allowed between linked sequences in clustering algorithm. A higher number will result in tighter clusters.",
207    )
208
209
210class ClusteredBatches(BaseModel):
211    """Configuration for clustered template selection. Clusters are sorted by the assay in `sort_parameters`
212    and the top values are passed for selection in batches.
213    """
214
215    assayed_sequences: (
216        list
217        | Annotated[
218            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
219            Field(discriminator="kind"),
220        ]
221    ) = Field(description="Table containing assayed sequences. Everything outside of this set will not be selected.")
222    sort_parameters: SortParameters
223    filter_parameters: FilterParameters | None = Field(default=None, description="If `None`, no filtering is applied.")
224    cluster_ids_or_parameters: (
225        list
226        | Annotated[
227            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
228            Field(discriminator="kind"),
229        ]
230        | ClusterParameters
231    ) = Field(
232        default_factory=ClusterParameters,
233        description="Either a table mapping sequences to cluster ids or parameters for running clustering.",
234    )
235
236
237class DatabaseSearch(BaseModel):
238    type_: Literal["DATABASE_SEARCH"] = Field(default="DATABASE_SEARCH", alias="type")
239    seed_domains: list[Annotated[str, Field(min_length=1, max_length=4000)]] = Field(
240        description="A list of domain sequences to search for.", min_length=1, max_length=100
241    )
242    database: ProteinDatabase = Field(
243        default=ProteinDatabase.UNIREF_30, description="Type of the protein database to use."
244    )
245
246
247class DiversifyParameters(ParametersBase):
248    """Parameters for the diversify task."""
249
250    task_type: ClassVar[str] = "diversify/v1"
251
252    template_sequence: DiversifyTemplateInputs = Field(description="The template amino acid sequence to diversify.")
253    homologs: (
254        list
255        | Annotated[
256            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
257            Field(discriminator="kind"),
258        ]
259    ) = Field(
260        description="Table of homologs to use for generator training. In the simplest case, this is just the result of a multiple sequence alignment (MSA) against the sequence to optimize."
261    )
262    domain_features: dict[
263        Annotated[str, Field(min_length=1, max_length=4000)],
264        Annotated[
265            list[DomainFeatureItem],
266            Field(
267                description="List of domain annotations. A residue may be assigned zero or one annotations. Ranges which are not annotated are considered reliable, i.e. the MSA in those ranges is used to infer conservation-based features."
268            ),
269        ],
270    ] = Field(
271        default_factory=dict,
272        description="Properties for each domain. A domain is one or more subsequences of the protein sequence, for example the active site of an enzyme, CDRs of an antibody, heavy and light chains of an scFv, etc. ",
273        max_length=100,
274    )
275    protein_model_type: ProteinModelType = Field(
276        default=ProteinModelType.DEFAULT,
277        description="Specifies the type of base model to use for training: use `DEFAULT` for single-chain proteins and `ANTIBODY` for multi-chain antibodies.",
278    )
279
280
281class DiversifyResult(ResultBase):
282    """Result of the diversify task."""
283
284    selected_sequences: TableResult = Field(description="Table containing the selected sequences.")
285    generated_sequences: TableResult = Field(description="Table containing the generated sequences.")
286    template_metadata: TemplateMetadata = Field(description="Metadata about templates.")
287    report: ArtifactParam = Field(description="A report on the result of the task.")
288
289
290class DiversifyTemplateInputs(BaseModel):
291    """Configuration for template-based sequence generation."""
292
293    template_id: str = Field(
294        description="Identifier for the template used to identify the source of generated sequences in the output."
295    )
296    sequence: str = Field(description="The template amino acid sequence", min_length=1, max_length=4000)
297    num_results: int = Field(
298        description="Number of sequences to generate for the final plate from the template sequence."
299    )
300    min_mutations: int = Field(default=1, description="Minimum number of mutations per sequence")
301    max_mutations: int = Field(default=4, description="Maximum number of mutations per sequence")
302    blocked_aas: list[BlockedAAItem] = Field(
303        default_factory=list,
304        description="List of blocked amino acids which may not be mutated to, with their respective ranges.",
305        max_length=500,
306    )
307    blocked_motifs: list[BlockedMotifItem] = Field(
308        default_factory=list,
309        description="List of blocked amino acid motifs with their respective ranges.",
310        max_length=1000,
311    )
312
313
314class DomainFeatureItem(BaseModel):
315    annotation: FeatureAnnotation = Field(
316        description="Annotation for the residues. `UNRELIABLE`: do not look at the MSA to infer conservation-based features (e.g. CDRs) for the protein `IGNORE`: do not change this position at all (e.g. for linkers) and do not use to infer any conservation-based features."
317    )
318    ranges: list[tuple[int, int]] = Field(
319        description="List of ranges where the annotation applies. 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) blocks positions 0 to 4."
320    )
321
322
323class EngineerParameters(ParametersBase):
324    task_type: ClassVar[str] = "engineer/v1"
325
326    dataset: (
327        list
328        | Annotated[
329            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
330            Field(discriminator="kind"),
331        ]
332    ) = Field(description="The table of assayed data, to be used for exploitation.")
333    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries", max_length=50)
334    primary_objective: PrimaryObjective = Field(description="The primary objective the samplers are trained for.")
335    constraints: list[Annotated[AbsoluteConstraint | RelativeConstraint, Field(discriminator="type_")]] = Field(
336        default_factory=list, description="List of assay constraints", max_length=50
337    )
338    samplers: list[ArtifactParam] = Field(description="The sampling models used to generate candidate sequences.")
339    scorer: ArtifactParam = Field(description="Scoring model used to predict assay values for candidate sequences.")
340    custom_predictors: list[CustomPredictorParam] = Field(
341        default_factory=list,
342        description="Specifications of custom predictors. If present, absolute constraints can be included to filter based on these predictors.",
343        max_length=10,
344    )
345    template_sequences: list[TemplateInputs] = Field(description="Initial sequences used as templates for generation")
346
347
348class EngineerResult(ResultBase):
349    selected_sequences: TableResult = Field(description="Table containing the selected sequences.")
350    engineered_sequences: TableResult = Field(description="Table containing the generated sequences.")
351    num_evaluated_seqs: int = Field(
352        description="Number of sequences evaluated in silico (only the ones sent to the predictor)."
353    )
354    num_generated_seqs: int = Field(
355        description="Number of sequences generated in silico (this is > num_evaluated_seqs, as it includes generated sequences that were not sent to the predictor, due to lower likelihood)."
356    )
357    new_mutation_ratio: float = Field(
358        description="Number of mutations in `selected_candidates` that do not appear in train/test datasets."
359    )
360    template_metadata: list[TemplateMetadata] = Field(description="Metadata about templates.")
361    report: ArtifactParam = Field(description="A report on the result of the task.")
362
363
364class FilterParameters(BaseModel):
365    """Parameters for filtering clusters based on assay values."""
366
367    objective: PrimaryObjective = Field(
368        description="The objective used for filtering clusters. Assay id must match one of the assay ids in the input dataset."
369    )
370    hit_redundancy: int = Field(
371        default=5,
372        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.',
373    )
374
375
376class Homologs(BaseModel):
377    type_: Literal["HOMOLOGS"] = Field(default="HOMOLOGS", alias="type")
378    homologs: (
379        list
380        | Annotated[
381            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
382            Field(discriminator="kind"),
383        ]
384    ) = Field(description="Precomputed homologs table.")
385
386
387class LabeledSelectionParameters(BaseModel):
388    """Configuration for down-selection based on labeled data. Selects the sequences with the highest labels."""
389
390    assayed_sequences: (
391        list
392        | Annotated[
393            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
394            Field(discriminator="kind"),
395        ]
396    ) = Field(
397        description="Table containing sequences with assays to use for down-selection. Everything outside of this set will not be selected."
398    )
399
400
401class PrimaryObjective(BaseModel):
402    assay_id: str = Field(description="The ID of the assay to be optimized.")
403    direction: OptimizationDirection = Field(description="Whether the maximize or minimize the assay value.")
404
405
406class RelativeConstraint(BaseModel):
407    type_: Literal["RELATIVE_CONSTRAINT"] = Field(default="RELATIVE_CONSTRAINT", alias="type")
408    assay_id: str = Field(description="The ID of the constrained assay.")
409    direction: OptimizationDirection = Field(
410        description="Whether the assay value of new sequences should be above (MAXIMIZE) or below (MINIMIZE) the assay value of the `relative_to` sequence."
411    )
412    relative_to: str = Field(
413        description="The sequence relative to which new sequences are constrained.", min_length=1, max_length=4000
414    )
415
416
417class ScoredSelectionParameters(BaseModel):
418    """Configuration for scorer-based down-selection."""
419
420    scorer: ArtifactParam
421    hit_redundancy: int = Field(
422        default=10,
423        description="Number of top candidates for which to optimize the average score. A higher number means less diversity.",
424    )
425    max_mutation_frequency: float | None = Field(
426        default=None,
427        description="Maximum mutation frequency in the set of selected sequences. If None, the algorithm will pick sensible defaults.",
428    )
429
430
431class SearchParameters(ParametersBase):
432    task_type: ClassVar[str] = "search/v1"
433
434    source: Homologs | DatabaseSearch = Field(
435        description="Precomputed homologs table or seed sequences to search for. If a table of precomputed sequence homologs is provided, the entries in the `seed_domain` column of this table may be used for sequence feature computation. For example, if the sequence to be optimized is an scFv, the `seed_domain` column would contain the heavy and light chain scaffold sequences of the scFv. These 2 scaffold sequences will be used to search for homologs and infer features such as the CDRs.",
436        discriminator="type_",
437    )
438    antibody_sequence_features: bool = Field(
439        default=False, description="Whether to compute sequence features. Only supported for antibodies."
440    )
441
442
443class SearchResult(ResultBase):
444    domain_features: dict[
445        Annotated[str, Field(min_length=1, max_length=4000)],
446        Annotated[
447            list[DomainFeatureItem],
448            Field(
449                description="List of domain annotations. A residue may be assigned zero or one annotations. Ranges which are not annotated are considered reliable, i.e. the MSA in those ranges is used to infer conservation-based features."
450            ),
451        ],
452    ] = Field(
453        default_factory=dict,
454        description="Mapping of protein sequences to their respective domain features. Each range in the domain features must be a valid index into the corresponding protein sequence.",
455        max_length=100,
456    )
457    homologs: TableResult
458
459
460class SelectParameters(ParametersBase):
461    """Select a set of sequences from a larger set of candidates.
462
463    Selection can be based on a scorer or on labeled data. If `clustered_batches` is specified, sequences are passed for selection in batches of cluster representatives.
464    """
465
466    task_type: ClassVar[str] = "select/v1"
467
468    sequences: (
469        list
470        | Annotated[
471            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
472            Field(discriminator="kind"),
473        ]
474    ) = Field(description="The sequences to select from.")
475    objective: PrimaryObjective = Field(
476        description="The primary objective to optimize during selection. This is used for scoring sequences."
477    )
478    constraints: list[Annotated[AbsoluteConstraint | RelativeConstraint, Field(discriminator="type_")]] = Field(
479        default_factory=list, description="Constraints that must be satisfied by the selected sequences.", max_length=50
480    )
481    selection_method: ScoredSelectionParameters | LabeledSelectionParameters = Field(
482        description="Method of selection. Can either be based on a scorer or on hard labels."
483    )
484    clustered_batches: ClusteredBatches | None = Field(
485        default=None,
486        description="Determines the clusters and representatives to use for selection. If `None`, all sequences are passed for selection at once.",
487    )
488    max_num_sequences_to_select: int = Field(description="No more than this many sequences will be selected.")
489    min_num_sequences_to_select: int = Field(
490        default=0, description="Warning will be raised if fewer than this many sequences are selected."
491    )
492
493
494class SelectResults(ResultBase):
495    selected_sequences: TableResult = Field(description="Table containing the selected sequences.")
496
497
498class SortParameters(BaseModel):
499    """Parameters for sorting clusters based on assay values."""
500
501    objective: PrimaryObjective = Field(
502        description="The objective used for sorting clusters. Assay id must match one of the assay ids in the input dataset. Sorting direction is determined by `objective.direction`."
503    )
504
505
506class TemplateInputs(BaseModel):
507    """Configuration for template-based sequence generation."""
508
509    template_id: str = Field(
510        description="Identifier for the template used to identify the source of generated sequences in the output.",
511        min_length=1,
512        max_length=50,
513    )
514    sequence: str = Field(description="The template amino acid sequence", min_length=1, max_length=4000)
515    num_results: int = Field(
516        description="Number of sequences generated from this template to be added to the final plate."
517    )
518    min_mutations: int = Field(default=1, description="Minimum number of mutations per sequence.")
519    max_mutations: int = Field(default=8, description="Maximum number of mutations per sequence.")
520    hit_redundancy: int = Field(
521        default=10,
522        description="Number of top candidates for which to optimize the average score. A higher number means less diversity.",
523    )
524    blocked_aas: list[BlockedAAItem] = Field(
525        default_factory=list,
526        description="List of blocked amino acids which may not be mutated to, with their respective ranges.",
527        max_length=500,
528    )
529    blocked_motifs: list[BlockedMotifItem] = Field(
530        default_factory=list,
531        description="List of blocked amino acid motifs with their respective ranges.",
532        max_length=1000,
533    )
534    blocked_regexps: list[str] | None = Field(
535        default_factory=list, description="Regular expressions that the sequence cannot match"
536    )
537    must_match_regexps: list[str] | None = Field(
538        default_factory=list,
539        description="Generated sequences must match all of the given regular expressions. Can be used to enforce certain motifs.",
540    )
541
542
543class TemplateMetadata(BaseModel):
544    template_id: str = Field(
545        description="Identifier for the template used to identify the source of generated sequences.",
546        min_length=1,
547        max_length=50,
548    )
549    sequence: str = Field(
550        description="Amino acid template sequence used for generation.", min_length=1, max_length=4000
551    )
552    blocked_aas: list[BlockedAAItem] = Field(
553        default_factory=list,
554        description="List of blocked amino acids which may not be mutated to, with their respective ranges.",
555        max_length=500,
556    )
557    min_mutations: int
558    max_mutations: int
559    structure: ArtifactParam | None = Field(default=None)
560
561
562class TrainParameters(ParametersBase):
563    task_type: ClassVar[str] = "train/v1"
564
565    homologs: (
566        list
567        | Annotated[
568            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
569            Field(discriminator="kind"),
570        ]
571    ) = Field(description="The table of homologous sequences to be used for training the base model.")
572    domain_features: dict[
573        Annotated[str, Field(min_length=1, max_length=4000)],
574        Annotated[
575            list[DomainFeatureItem],
576            Field(
577                description="List of domain annotations. A residue may be assigned zero or one annotations. Ranges which are not annotated are considered reliable, i.e. the MSA in those ranges is used to infer conservation-based features."
578            ),
579        ],
580    ] = Field(
581        default_factory=dict,
582        description="Mapping of protein sequences to their respective domain features. Each range in the domain features must be a valid index into the corresponding protein sequence.",
583        max_length=100,
584    )
585    protein_model_type: ProteinModelType = Field(
586        default=ProteinModelType.DEFAULT,
587        description="Specifies the type of base model to use for training: use `DEFAULT` for single-chain proteins and `ANTIBODY` for multi-chain antibodies.",
588    )
589    dataset: (
590        list
591        | Annotated[
592            TableSourceTable | TableSourceQuery | TableSourceTaskResult | TableSourceSnapshot | TableSourceLiteral,
593            Field(discriminator="kind"),
594        ]
595    ) = Field(description="The name of the input table to be used for model training.")
596    assays: list[Assay] = Field(default_factory=list, description="List of assay metadata entries", max_length=50)
597    primary_objective: PrimaryObjective = Field(
598        description="The primary objective on which the sampler model is conditioned. Should correspond to an assay in `dataset`."
599    )
600
601
602class TrainResult(ResultBase):
603    base_sampler: ArtifactParam = Field(
604        description="A sampler which attempts to produce plausible sequences as understood relative to the passed `homologs`. "
605    )
606    conditioned_sampler: ArtifactParam | None = Field(
607        description="A sampler which attempts to produce high quality sequences as understood relative to the passed `primary_objective`."
608    )
609    scorer: ArtifactParam = Field(description="The scoring model, trained on the assay data provided.")
610    prediction_data: TableResult = Field(description="Predicted assay values on all folds of the input dataset.")
611    generator_prediction_data: TableResult = Field(
612        description="Pseudo log-likelihoods of the sequences in the dataset under the generator(s)."
613    )
614    report: ArtifactParam = Field(description="A report on the result of the task.")