RITA_m / README.md
DanielHesslow's picture
Update README.md
b4bfffb
metadata
language: protein
tags:
  - protein
datasets:
  - uniref-100

RITA-M

RITA is a family of autoregressive protein models, developed by a collaboration of Lighton, the OATML group at Oxford, and the Debbie Marks Lab at Harvard.

Model #Params d_model layers lm loss uniref-100
Small 85M 768 12 2.31
Medium 300M 1024 24 2.01
Large 680M 1536 24 1.82
XLarge 1.2B 2048 24 1.70

For full results see our preprint: https://arxiv.org/abs/2205.05789

Usage

Instantiate a model like so:

from transformers import AutoModel, AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("lightonai/RITA_m, trust_remote_code=True")
tokenizer = AutoTokenizer.from_pretrained("lightonai/RITA_m")

for generation we support pipelines:

from transformers import pipeline
rita_gen = pipeline('text-generation', model=model, tokenizer=tokenizer)
sequences = rita_gen("MAB", max_length=20, do_sample=True, top_k=950, repetition_penalty=1.2, 
                     num_return_sequences=2, eos_token_id=2)
for seq in sequences:
    print(f"seq: {seq['generated_text'].replace(' ', '')}")

How to cite

@misc{RITA2022,
  doi = {10.48550/ARXIV.2205.05789},
  url = {https://arxiv.org/abs/2205.05789},
  author = {Hesslow, Daniel and Zanichelli, Niccoló and Notin, Pascal and Poli, Iacopo and Marks, Debora},
  title = {RITA: a Study on Scaling Up Generative Protein Sequence Models},
  publisher = {arXiv},
  year = {2022},
}