--- license: mit datasets: - chandar-lab/UR100P language: - en tags: - biology --- ## AMPLIFY AMPLIFY is an efficient, state-of-the-art protein language model pre-trained using masked language modeling on UniRef100, OAS, and SCOP ([UR100P](https://huggingface.co/datasets/chandar-lab/UR100P)). AMPLIFY can generate residue and protein embeddings, suggest mutations, differentiate disordered proteins from non-protein sequences, and much more. AMPLIFY is available in two sizes, 120M and 350M parameters, with the `_base` models not extended beyond 512 residues (Stage 1). The model architecture and pre-training procedure are detailed below. For more details, please refer to the [accompanying paper](https://www.biorxiv.org/content/10.1101/2024.09.23.614603v1). - [`AMPLIFY_350M`](https://huggingface.co/chandar-lab/AMPLIFY_350M) - [`AMPLIFY_350M_base`](https://huggingface.co/chandar-lab/AMPLIFY_350M_base) - [`AMPLIFY_120M`](https://huggingface.co/chandar-lab/AMPLIFY_120M) - [`AMPLIFY_120M_base`](https://huggingface.co/chandar-lab/AMPLIFY_120M_base) ### Model Descritpion | | AMPLIFY 120M | AMPLIFY 350M | | :----------------------------- | -----------: | -----------: | | `hidden-size` | 640 | 960 | | `num-hidden-layers` | 24 | 32 | | `num-attention-heads` | 10 | 15 | | `intermediate-size` | 2560 | 3840 | | `max-position-embeddings` | 2048 | 2048 | | `vocab-size` | 27 | 27 | | `rope-theta` | 10000 | 10000 | | `dropout-prob` | 0 | 0 | | `embedding-init-range` | 0.02 | 0.02 | | `norm-eps` | 1.0e-05 | 1.0e-05 | | `hidden-act` | swiglu | swiglu | | `pre-activation-layer-norm` | true | true | | `layer-norm-after-embedding` | false | false | | `layer-norm-before-last-layer` | true | true | | `rms-norm` | true | true | | `ffn-bias` | false | false | | `attn-bias` | false | false | ### Training Descritpion | | Stage 1 | Stage 2 | | :------------------ | ----------: | ---------------------------: | | `dataset` | UR100P | UR100P | | `max-steps` | 1000000 | 25000 (120M) or 50000 (350M) | | `max-length` | 512 | 2048 | | `optimizer` | adamw | adamw | | `lr` | 0.001 | 0.001 | | `betas` | (0.9, 0.95) | (0.9, 0.95) | | `eps` | 1.0e-08 | 1.0e-08 | | `weight-decay` | 0.01 | 0.01 | | `scheduler` | cosinedecay | none | | `warmup-steps` | 1,000 | none | | `final-step` | 900,000 | none | | `warmup-steps` | 1,000 | none | | `gradient-clipping` | 1.0 | 1.0 | | `tf32` | true | true | | `mixed-precision` | bf16 | bf16 | | `padding` | max-length | max-length | | `random-truncate` | true | true | | `mask-probability` | 0.15 | 0.15 | | `total-batch-size` | 4096 | 4096 | | `deepspeed` | true | true | | `zero-stage` | 3 | 3 | ## Get Started ```python from transformers import AutoModel from transformers import AutoTokenizer from datasets import load_dataset # Load AMPLIFY and tokenizer model = AutoModel.from_pretrained("chandar-lab/AMPLIFY_350M", trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained("chandar-lab/AMPLIFY_350M", trust_remote_code=True) # Move the model to GPU (required due to Flash Attention) model = model.to("cuda") # Load the UniProt validation set dataset = load_dataset("chandar-lab/UR100P", data_dir="UniProt", split="test") for sample in dataset: # Protein print("Sample: ", sample["name"], sample["sequence"]) # Tokenize the protein input = tokenizer.encode(sample["sequence"], return_tensors="pt") print("Input: ", input) # Move to the GPU and make a prediction input = input.to("cuda") output = model(input) print("Output: ", output) break ``` ## Citations If you find the models useful in your research, we ask that you cite the paper: ```bibtex @article{Fournier2024.09.23.614603, title = {Protein Language Models: Is Scaling Necessary?}, author = {Fournier, Quentin and Vernon, Robert M. and van der Sloot, Almer and Schulz, Benjamin and Chandar, Sarath and Langmead, Christopher James}, year = {2024}, journal = {bioRxiv}, publisher = {Cold Spring Harbor Laboratory}, doi = {10.1101/2024.09.23.614603}, url = {https://www.biorxiv.org/content/early/2024/09/23/2024.09.23.614603}, elocation-id = {2024.09.23.614603}, eprint = {https://www.biorxiv.org/content/early/2024/09/23/2024.09.23.614603.full.pdf} } ```