Akselerer prediksjon av proteinstruktur med ESMFold-språkmodellen på Amazon SageMaker

Proteiner driver mange biologiske prosesser, som enzymaktivitet, molekylær transport og cellulær støtte. Den tredimensjonale strukturen til et protein gir innsikt i dets funksjon og hvordan det interagerer med andre biomolekyler. Eksperimentelle metoder for å bestemme proteinstruktur, som røntgenkrystallografi og NMR-spektroskopi, er kostbare og tidkrevende.

I motsetning til dette kan nylig utviklede beregningsmetoder raskt og nøyaktig forutsi strukturen til et protein fra dets aminosyresekvens. Disse metodene er kritiske for proteiner som er vanskelige å studere eksperimentelt, for eksempel membranproteiner, målene for mange medikamenter. Et kjent eksempel på dette er AlphaFold, en dyp læringsbasert algoritme kjent for sine nøyaktige spådommer.

ESMFold er en annen svært nøyaktig, dyp læringsbasert metode utviklet for å forutsi proteinstruktur fra aminosyresekvensen. ESMFold bruker en stor proteinspråkmodell (pLM) som ryggrad og opererer ende til ende. I motsetning til AlphaFold2, trenger den ikke et oppslag eller Multiple Sequence Alignment (MSA) trinn, og er heller ikke avhengig av eksterne databaser for å generere spådommer. I stedet trente utviklingsteamet modellen på millioner av proteinsekvenser fra UniRef. Under trening utviklet modellen oppmerksomhetsmønstre som elegant representerer de evolusjonære interaksjonene mellom aminosyrer i sekvensen. Denne bruken av en pLM i stedet for en MSA muliggjør opptil 60 ganger raskere prediksjonstider enn andre toppmoderne modeller.

I dette innlegget bruker vi den ferdigtrente ESMFold-modellen fra Hugging Face with Amazon SageMaker å forutsi tungkjedestrukturen til trastuzumaben monoklonalt antistoff først utviklet av Genentech for behandling av HER2-positiv brystkreft. Å raskt forutsi strukturen til dette proteinet kan være nyttig hvis forskere ønsket å teste effekten av sekvensmodifikasjoner. Dette kan potensielt føre til forbedret pasientoverlevelse eller færre bivirkninger.

Dette innlegget gir et eksempel på en Jupyter-notatbok og relaterte skript i det følgende GitHub repository.

Forutsetninger

Vi anbefaler å kjøre dette eksemplet i en Amazon SageMaker Studio bærbare kjører PyTorch 1.13 Python 3.9 CPU-optimalisert bilde på en ml.r5.xlarge-forekomsttype.

Visualiser den eksperimentelle strukturen til trastuzumab

Til å begynne med bruker vi biopython bibliotek og et hjelpeskript for å laste ned trastuzumab-strukturen fra RCSB Protein Data Bank:

from Bio.PDB import PDBList, MMCIFParser
from prothelpers.structure import atoms_to_pdb target_id = "1N8Z"
pdbl = PDBList()
filename = pdbl.retrieve_pdb_file(target_id, pdir="data")
parser = MMCIFParser()
structure = parser.get_structure(target_id, filename)
pdb_string = atoms_to_pdb(structure)

Deretter bruker vi py3Dmol bibliotek for å visualisere strukturen som en interaktiv 3D-visualisering:

view = py3Dmol.view()
view.addModel(pdb_string)
view.setStyle({'chain':'A'},{"cartoon": {'color': 'orange'}})
view.setStyle({'chain':'B'},{"cartoon": {'color': 'blue'}})
view.setStyle({'chain':'C'},{"cartoon": {'color': 'green'}})
view.show()

Følgende figur representerer 3D-proteinstrukturen 1N8Z fra Protein Data Bank (PDB). I dette bildet vises den lette kjeden av trastuzumab i oransje, den tunge kjeden er blå (med den variable regionen i lyseblått), og HER2-antigenet er grønt.

Vi vil først bruke ESMFold til å forutsi strukturen til den tunge kjeden (kjede B) fra dens aminosyresekvens. Deretter vil vi sammenligne prediksjonen med den eksperimentelt bestemte strukturen vist ovenfor.

Forutsi trastuzumab tungkjedestruktur fra dens sekvens ved å bruke ESMFold

La oss bruke ESMFold-modellen til å forutsi strukturen til den tunge kjeden og sammenligne den med det eksperimentelle resultatet. Til å begynne med bruker vi et forhåndsbygd notatbokmiljø i Studio som kommer med flere viktige biblioteker, som PyTorch, forhåndsinstallert. Selv om vi kunne bruke en akselerert forekomsttype for å forbedre ytelsen til den bærbare analysen vår, bruker vi i stedet en ikke-akselerert forekomst og kjører ESMFold-prediksjonen på en CPU.

Først laster vi den ferdigtrente ESMFold-modellen og tokenizeren fra Klemmer Face Hub:

from transformers import AutoTokenizer, EsmForProteinFolding tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)

Deretter kopierer vi modellen til enheten vår (CPU i dette tilfellet) og setter noen modellparametere:

device = torch.device("cpu")
model.esm = model.esm.float()
model = model.to(device)
model.trunk.set_chunk_size(64)

For å forberede proteinsekvensen for analyse, må vi tokenisere den. Dette oversetter aminosyresymbolene (EVQLV...) til et numerisk format som ESMFold-modellen kan forstå (6,19,5,10,19,...):

tokenized_input = tokenizer([experimental_sequence], return_tensors="pt", add_special_tokens=False)["input_ids"]
tokenized_input = tokenized_input.to(device)

Deretter kopierer vi den tokeniserte inngangen til modusen, lager en prediksjon og lagrer resultatet i en fil:

with torch.no_grad():
notebook_prediction = model.infer_pdb(experimental_sequence)
with open("data/prediction.pdb", "w") as f:
f.write(notebook_prediction)

Dette tar omtrent 3 minutter på en ikke-akselerert forekomsttype, som en r5.

Vi kan sjekke nøyaktigheten til ESMFold-prediksjonen ved å sammenligne den med den eksperimentelle strukturen. Vi gjør dette ved å bruke US-Align verktøy utviklet av Zhang Lab ved University of Michigan:

from prothelpers.usalign import tmscore tmscore("data/prediction.pdb", "data/experimental.pdb", pymol="data/superimposed")

De mal modellering poengsum (TM-score) er en beregning for å vurdere likheten mellom proteinstrukturer. En poengsum på 1.0 indikerer en perfekt match. Poeng over 0.7 indikerer at proteiner deler samme ryggradsstruktur. Poeng over 0.9 indikerer at proteinene er funksjonelt utskiftbare for nedstrøms bruk. I vårt tilfelle med å oppnå TM-score 0.802, vil ESMFold-prediksjonen sannsynligvis være passende for applikasjoner som strukturscoring eller ligandbindingseksperimenter, men kan ikke være egnet for brukstilfeller som molekylær erstatning som krever ekstremt høy nøyaktighet.

Vi kan validere dette resultatet ved å visualisere de justerte strukturene. De to strukturene viser en høy, men ikke perfekt, grad av overlapping. Proteinstrukturspådommer er et felt i rask utvikling, og mange forskerteam utvikler stadig mer nøyaktige algoritmer!

Distribuer ESMFold som et SageMaker-slutningsendepunkt

Å kjøre modellslutninger i en bærbar PC er greit for eksperimentering, men hva om du trenger å integrere modellen med en applikasjon? Eller en MLOps-rørledning? I dette tilfellet er et bedre alternativ å distribuere modellen din som et slutningsendepunkt. I det følgende eksempelet vil vi distribuere ESMFold som et SageMaker-endepunkt for sanntidsslutning på en akselerert forekomst. SageMaker sanntidsendepunkter gir en skalerbar, kostnadseffektiv og sikker måte å distribuere og være vert for maskinlæringsmodeller (ML). Med automatisk skalering kan du justere antall forekomster som kjører endepunktet for å møte kravene til applikasjonen din, optimalisere kostnadene og sikre høy tilgjengelighet.

Den ferdigbygde SageMaker beholder for Hugging Face gjør det enkelt å distribuere dyplæringsmodeller for vanlige oppgaver. For nye brukstilfeller som prediksjon av proteinstruktur, må vi imidlertid definere en skikk inference.py skript for å laste modellen, kjøre prediksjonen og formatere utdataene. Dette skriptet inneholder mye av den samme koden som vi brukte i notatboken vår. Vi lager også en requirements.txt fil for å definere noen Python-avhengigheter som endepunktet vårt skal bruke. Du kan se filene vi opprettet i GitHub repository.

I den følgende figuren er de eksperimentelle (blå) og forutsagte (røde) strukturene til den tunge kjeden av trastuzumab svært like, men ikke identiske.

Etter at vi har opprettet de nødvendige filene i code katalogen, distribuerer vi modellen vår ved å bruke SageMaker HuggingFaceModel klasse. Dette bruker en forhåndsbygd beholder for å forenkle prosessen med å distribuere Hugging Face-modeller til SageMaker. Merk at det kan ta 10 minutter eller mer å opprette endepunktet, avhengig av tilgjengeligheten til ml.g4dn instanstyper i vår region.

from sagemaker.huggingface import HuggingFaceModel
from datetime import datetime huggingface_model = HuggingFaceModel(
model_data = model_artifact_s3_uri, # Previously staged in S3
name = f"emsfold-v1-model-" + datetime.now().strftime("%Y%m%d%s"),
transformers_version='4.17',
pytorch_version='1.10',
py_version='py38',
role=role,
source_dir = "code",
entry_point = "inference.py"
) rt_predictor = huggingface_model.deploy(
initial_instance_count = 1,
instance_type="ml.g4dn.2xlarge",
endpoint_name=f"my-esmfold-endpoint",
serializer = sagemaker.serializers.JSONSerializer(),
deserializer = sagemaker.deserializers.JSONDeserializer()
)

Når endepunktsimplementeringen er fullført, kan vi sende inn proteinsekvensen på nytt og vise de første radene i prediksjonen:

endpoint_prediction = rt_predictor.predict(experimental_sequence)[0]
print(endpoint_prediction[:900])

Fordi vi distribuerte endepunktet vårt til en akselerert forekomst, bør prediksjonen bare ta noen få sekunder. Hver rad i resultatet tilsvarer et enkelt atom og inkluderer aminosyreidentiteten, tre romlige koordinater og en pLDDT-score som representerer prediksjonstilliten på det stedet.

Ved å bruke samme metode som før ser vi at notisboken og endepunktsprediksjonene er identiske.

Som observert i den følgende figuren, viser ESMFold-spådommene generert i den bærbare datamaskinen (rød) og ved endepunktet (blått) perfekt justering.

Rydd opp

For å unngå ytterligere gebyrer sletter vi slutningsendepunktet og testdataene våre:

rt_predictor.delete_endpoint()
bucket = boto_session.resource("s3").Bucket(bucket)
bucket.objects.filter(Prefix=prefix).delete()
os.system("rm -rf data obsolete code")

Oppsummering

Beregningsbasert proteinstrukturprediksjon er et kritisk verktøy for å forstå funksjonen til proteiner. I tillegg til grunnforskning har algoritmer som AlphaFold og ESMFold mange bruksområder innen medisin og bioteknologi. Den strukturelle innsikten generert av disse modellene hjelper oss å bedre forstå hvordan biomolekyler samhandler. Dette kan da føre til bedre diagnostiske verktøy og terapier for pasienter.

I dette innlegget viser vi hvordan du distribuerer ESMFold-proteinspråkmodellen fra Hugging Face Hub som et skalerbart sluttpunkt ved bruk av SageMaker. For mer informasjon om distribusjon av Hugging Face-modeller på SageMaker, se Bruk Hugging Face med Amazon SageMaker. Du kan også finne flere proteinvitenskapelige eksempler i Fantastisk proteinanalyse på AWS GitHub repo. Legg igjen en kommentar hvis det er noen andre eksempler du vil se!

Om forfatterne

Brian lojal er senior AI/ML Solutions Architect i Global Healthcare and Life Sciences-teamet hos Amazon Web Services. Han har mer enn 17 års erfaring innen bioteknologi og maskinlæring, og brenner for å hjelpe kunder med å løse genomiske og proteomiske utfordringer. På fritiden liker han å lage mat og spise sammen med venner og familie.

Shamika Ariyawansa er en AI/ML-spesialistløsningsarkitekt i Global Healthcare and Life Sciences-teamet hos Amazon Web Services. Han jobber lidenskapelig med kunder for å akselerere deres AI- og ML-adopsjon ved å gi teknisk veiledning og hjelpe dem med å innovere og bygge sikre skyløsninger på AWS. Utenom jobben elsker han ski og terrengkjøring.

Yanjun Qi er Senior Applied Science Manager ved AWS Machine Learning Solution Lab. Hun innoverer og bruker maskinlæring for å hjelpe AWS-kunder med å få fart på AI og skyadopsjon.

• SEO-drevet innhold og PR-distribusjon. Bli forsterket i dag.
• PlatoAiStream. Web3 Data Intelligence. Kunnskap forsterket. Tilgang her.
• Minting the Future med Adryenn Ashley. Tilgang her.
• Kjøp og selg aksjer i PRE-IPO-selskaper med PREIPO®. Tilgang her.
• kilde: https://aws.amazon.com/blogs/machine-learning/accelerate-protein-structure-prediction-with-the-esmfold-language-model-on-amazon-sagemaker/