Trias-Small: a mammalian variant of Trias for codon generation

Trias-Small is a variant of Trias, the encoder-decoder model that reverse-translates protein sequences into codon sequences. Unlike Trias, which was trained on vertebrate mRNA sequences across many species, Trias-Small was trained only on mammalian coding sequences.

For full setup, decoding options (greedy / beam search), and citation, see the main lareaulab/Trias model card. This card only documents what differs.

Key difference from Trias

Trias-Small is trained only on mammalian sequences and does not take an input species tag — there is no >>{species}<< prefix.

Reverse Translation

Trias-Small uses the custom TriasTokenizer, so the trias package must be installed:

pip install git+https://github.com/lareaulab/Trias.git

import torch
from huggingface_hub import hf_hub_download
from transformers import BartForConditionalGeneration
from trias.tokenizer import TriasTokenizer

repo = "lareaulab/Trias-Small"

# The tokenizer is built directly from the vocab file
vocab_file = hf_hub_download(repo, "vocab.json")
tokenizer = TriasTokenizer(vocab_file, separate_vocabs=False, model_max_length=2048)
model = BartForConditionalGeneration.from_pretrained(repo).eval()

# No species tag — just the protein sequence (end with '*' for the stop)
protein_sequence = "MTEITAAMVKELRESTGAGMMDCKNALSETQ*"

input_ids = tokenizer.encode(protein_sequence, return_tensors="pt")
outputs = model.generate(input_ids, max_length=tokenizer.model_max_length)
codon_sequence = tokenizer.decode(outputs[0], skip_special_tokens=True)

print("Codon sequence:", codon_sequence)

Beam search works the same way as in the main Trias model (pass num_beams=... to model.generate).