Runs the full inference pipeline to predict structures for each seed.
(
fold_input: folding_input.Input,
model_runner: ModelRunner,
*,
buckets: Sequence[int] | None = None,
ref_max_modified_date: datetime.date | None = None,
conformer_max_iterations: int | None = None,
resolve_msa_overlaps: bool = True,
fix_standalone_glycans: bool = False,
)
| 526 | |
| 527 | |
| 528 | def predict_structure( |
| 529 | fold_input: folding_input.Input, |
| 530 | model_runner: ModelRunner, |
| 531 | *, |
| 532 | buckets: Sequence[int] | None = None, |
| 533 | ref_max_modified_date: datetime.date | None = None, |
| 534 | conformer_max_iterations: int | None = None, |
| 535 | resolve_msa_overlaps: bool = True, |
| 536 | fix_standalone_glycans: bool = False, |
| 537 | ) -> Sequence[ResultsForSeed]: |
| 538 | """Runs the full inference pipeline to predict structures for each seed.""" |
| 539 | |
| 540 | print(f'Featurising data with {len(fold_input.rng_seeds)} seed(s)...') |
| 541 | featurisation_start_time = time.time() |
| 542 | ccd = chemical_components.Ccd(user_ccd=fold_input.user_ccd) |
| 543 | featurised_examples = featurisation.featurise_input( |
| 544 | fold_input=fold_input, |
| 545 | buckets=buckets, |
| 546 | ccd=ccd, |
| 547 | verbose=True, |
| 548 | ref_max_modified_date=ref_max_modified_date, |
| 549 | conformer_max_iterations=conformer_max_iterations, |
| 550 | resolve_msa_overlaps=resolve_msa_overlaps, |
| 551 | fix_standalone_glycans=fix_standalone_glycans, |
| 552 | ) |
| 553 | print( |
| 554 | f'Featurising data with {len(fold_input.rng_seeds)} seed(s) took' |
| 555 | f' {time.time() - featurisation_start_time:.2f} seconds.' |
| 556 | ) |
| 557 | print( |
| 558 | 'Running model inference and extracting output structure samples with' |
| 559 | f' {len(fold_input.rng_seeds)} seed(s)...' |
| 560 | ) |
| 561 | all_inference_start_time = time.time() |
| 562 | all_inference_results = [] |
| 563 | for seed, example in zip(fold_input.rng_seeds, featurised_examples): |
| 564 | print(f'Running model inference with seed {seed}...') |
| 565 | inference_start_time = time.time() |
| 566 | rng_key = jax.random.PRNGKey(seed) |
| 567 | result = model_runner.run_inference(example, rng_key) |
| 568 | print( |
| 569 | f'Running model inference with seed {seed} took' |
| 570 | f' {time.time() - inference_start_time:.2f} seconds.' |
| 571 | ) |
| 572 | print(f'Extracting inference results with seed {seed}...') |
| 573 | extract_structures = time.time() |
| 574 | inference_results = model_runner.extract_inference_results( |
| 575 | batch=example, result=result, target_name=fold_input.name |
| 576 | ) |
| 577 | num_tokens = len(inference_results[0].metadata['token_chain_ids']) |
| 578 | embeddings = model_runner.extract_embeddings( |
| 579 | result=result, num_tokens=num_tokens |
| 580 | ) |
| 581 | distogram = model_runner.extract_distogram( |
| 582 | result=result, num_tokens=num_tokens |
| 583 | ) |
| 584 | print( |
| 585 | f'Extracting {len(inference_results)} inference samples with' |
no test coverage detected