The ability to efficiently generate specific RNA sequences on demand has significant implications for both scientific research and therapeutic applications. In this context, we introduce RFamLlama, a conditional language model that is specifically optimized for generating RNA sequences across diverse families. This model was trained on RNA sequences representing over 4,000 distinct families, each augmented with control tags to denote the specific family. We have shown that the inclusion of family-specific tags substantially enhances the capabilities of our model in zero-shot fitness prediction of RNA molecules. Additionally, this model supports a conditional generation approach, allowing for the generation of RNA sequences by using Rfam IDs as input prompts, thereby eliminating the need for further functional-specific fine-tuning. Consequently, RFamLlama is poised to be an effective and widely applicable tool for the zero-shot fitness prediction and generation of RNA sequences, potentially pushing the boundaries of what can be achieved beyond natural evolutionary processes.
1# generation
2from transformers import LlamaForCausalLM, AutoTokenizer, pipeline
3import torch
4import sys
5
6model_url = "jinyuan22/RFamLlama-base"
7model = LlamaForCausalLM.from_pretrained(model_url, torch_dtype=torch.float16)
8tokenizer = AutoTokenizer.from_pretrained(model_url)
9
10device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
11
12pipe = pipeline("text-generation", model=model, device=device, tokenizer=tokenizer)
13
14tag = "RF00005"
15txt = f"<|bos|> <|tag_start|> {tag[2:]} <|tag_end|> <|5|> "
16all_outputs = []
17outputs = pipe(txt, num_return_sequences=10, max_new_tokens=300, repetition_penalty=1, top_p=1,temperature=1, do_sample=True)
18
19for i, output in enumerate(outputs):
20 seq = output["generated_text"]
21 seq = seq.split("<|5|>")[1].split("<|3|>")[0]
22 print(f">{i}\n{seq}")
