Update README.md

README.md

@@ -2,6 +2,8 @@
 license: mit
 tags:
 - biology
+- mrna design
+- codon optimization
 pipeline_tag: translation
 ---
 
@@ -19,32 +21,34 @@ Greedy decoding selects the most likely token at each step, it's faster and dete
 
 
 ```python
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, BartForConditionalGeneration
 
 # Load model and tokenizer from the Hub
-tokenizer = AutoTokenizer.from_pretrained("lareaulab/Trias", use_fast=False)
-model = AutoModelForSeq2SeqLM.from_pretrained("lareaulab/Trias")
+tokenizer = AutoTokenizer.from_pretrained("lareaulab/Trias", use_fast=True)
+model = BartForConditionalGeneration.from_pretrained("lareaulab/Trias")
 
 # Input sequence
 species = "Homo sapiens"
 protein_sequence = "MTEITAAMVKELRESTGAGMMDCKNALSETQ*"
-input_text = f"{species}: {protein_sequence}"
+input_seq = f">>{species}<< {protein_sequence}"
 
 # Tokenize
-inputs = tokenizer(input_text, return_tensors="pt")
+input_ids = tokenizer.encode(input_seq, return_tensors="pt")
 
 # Generate codon sequence (greedy)
-outputs = model.generate(**inputs)
+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 example 
 ```python
 outputs = model.generate(
-    **inputs,
+    input_ids,
     num_beams=5,
-    early_stopping=True
+    early_stopping=True,
+    max_length=tokenizer.model_max_length)
 )
 ```