Update README.md

README.md

@@ -38,12 +38,12 @@ import torch
 from transformers import AutoTokenizer, BigBirdForMaskedLM
 from CodonTransformer.CodonPrediction import predict_dna_sequence
 from CodonTransformer.CodonJupyter import format_model_output
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("adibvafa/CodonTransformer")
-model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer-base").to(DEVICE)
+model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer-base").to(device)
 
 
 # Set your input data
@@ -55,10 +55,11 @@ organism = "Escherichia coli general"
 output = predict_dna_sequence(
     protein=protein,
     organism=organism,
-    device=DEVICE,
+    device=device,
     tokenizer=tokenizer,
     model=model,
     attention_type="original_full",
+    deterministic=True
 )
 print(format_model_output(output))
 ```